URUMETRICS-CODE/import_data/utils.py

#!usr/bin/env python
# -*- coding: utf8 -*-

# -----------------------------------------------------------------------------
#   File: utils.py (as part of project URUMETRICS)
#   Created: 04/05/2022 15:47
#   Last Modified: 04/05/2022 15:47
# -----------------------------------------------------------------------------
#   Author: William N. Havard
#           Postdoctoral Researcher
#
#   Mail  : william.havard@ens.fr / william.havard@gmail.com
#  
#   Institution: ENS / Laboratoire de Sciences Cognitives et Psycholinguistique
#
# ------------------------------------------------------------------------------
#   Description: 
#       • 
# -----------------------------------------------------------------------------

import os


def get_raw_filename(fp):
    """
    Return the raw filename of file given its path (removes path and extension)
    :param fp: file path
    :type fp: str
    :return: raw filename
    :rtype: str
    """
    return os.path.basename(os.path.splitext(fp)[0])


def walk_dir(path, ext=[], return_full_path=True):
    """
    Recursively explore a directory and return all the files that end with a given extension. If no extension is
    specified, all the files are returned
    :param path: path to the directory to explore
    :type path: str
    :param ext: extensions that should be considered
    :type ext: list of str or str
    :param return_full_path: should the full path to the files be returned
    :type return_full_path: bool
    :return: list of files
    :rtype: list of str
    """
    if type(ext) == str: ext = [ext]

    files = []
    for p, d, f in os.walk(path):
        for ff in f:
            _, file_extension = os.path.splitext(ff)
            if not len(ext) or file_extension[1:] in ext:
                path_suffix = p.replace(os.path.commonprefix([path, p]), '').lstrip(os.sep)
                if return_full_path:
                    files.append(os.path.join(path, path_suffix, ff))
                else:
                    files.append(os.path.join(path_suffix, ff))
    return sorted(files)