get_files_structure.py

import json
import os


def set_nested_dict_value(d, keys, value):
    d_next = d
    for key in keys:
        if '_dirs' not in d_next:
            d_next['_dirs'] = dict()
        if key not in d_next['_dirs']:
            d_next['_dirs'][key] = dict()
        d_next = d_next['_dirs'][key]
    else:
        if '_files' not in d_next:
            d_next['_files'] = []
        d_next['_files'].append(value)


def get_folder_structure(root_dir, return_type='tree', follow_symlinks=False):
    """
    Recursively get folder structure.
    :param root_dir:  The path to the folder.
    :param return_type: tree or flatten. if flatten just return a list of files not considering the folder structure.
    :param follow_symlinks: flag to follow symlinks.
    :return:

    Example:
     :return_type: tree
        {
            "_dirs": {
                "folder1": {
                    "_dirs": {
                        "folder2": {
                            "_files": ["file1.txt", "file2.txt"]
                        }
                    }
                }
            }
        }

        :return_type: flatten
        [
            "folder1/folder2/file1.txt",
            "folder1/folder2/file2.txt"
        ]

    """
    assert return_type in ['tree', 'flatten'], "return_type must be either 'tree' or 'flat'"
    assert os.path.exists(root_dir), "Directory does not exist"

    if return_type == 'tree':
        files_structure_tr = dict()
        for root1, dirs1, files1 in os.walk(root_dir, followlinks=follow_symlinks):
            rep1 = root1.split(os.sep)[len(root_dir.split(os.sep)):]
            for name in files1:
                set_nested_dict_value(files_structure_tr, rep1, name)
        return files_structure_tr
    else:
        files_structure_flt = []
        for root, dirs, files in os.walk(root_dir, followlinks=follow_symlinks):
            rep = f"{os.sep}".join(root.split(os.sep)[len(root_dir.split(os.sep)):])
            files_structure_flt.extend([os.path.join(rep, file_name) for file_name in files])
        return files_structure_flt


if __name__ == '__main__':
    folder_structure = get_folder_structure('.', return_type='flatten')

    with open('st.json', 'w') as f:
        json.dump(folder_structure, f, indent=4)