diff --git a/src/tree_repr.py b/src/tree_repr.py new file mode 100644 index 0000000..b42c786 --- /dev/null +++ b/src/tree_repr.py @@ -0,0 +1,298 @@ +import os + +EMPTY_TREE_LINE_LENGTH = 160 +EMPTY_TREE_LINE = ' ' * EMPTY_TREE_LINE_LENGTH + +VERTICAL_CONNECTOR = '│' +### same length +CHILD_CONNECTOR = '├── ' +LAST_CHILD_CONNECTOR = '└── ' +### is the length of child connector and last_child_connector +LEVEL_INDENT = 4 +ITEMS_TO_EXCLUDE = ['__pycache__'] + +def mount_str_at(source_str, index, mount_str): + l = len(mount_str) + return source_str[:index] + mount_str + source_str[index+l:] + +def level2indent(l): + if not l: + return 0 + + return LEVEL_INDENT * l + +def produce_treeline(prec_seps, lsep): + + tree_line = EMPTY_TREE_LINE[:] + lseps = prec_seps[:] + lseps.append(lsep) + + for level_n, (sep, item) in enumerate(lseps): + tomount = sep + if item: + tomount = sep + item + tree_line = mount_str_at(tree_line, level2indent(level_n), tomount) + + return tree_line.rstrip() + +def nerd_tree_sort_dirfirst(startpath): + files = [] + dirs = [] + + for item in os.listdir(startpath): + if os.path.isdir(startpath + item): + dirs.append(item) + continue + files.append(item) + + return dirs + files + +def nerd_tree_sort(startpath, dirfirst=True): + ''' + it returns the sorted list of items starting from startpath + ''' + + # dirs = [] + # files = [] + + if dirfirst: + return nerd_tree_sort_dirfirst(startpath) + + return os.listdir(startpath) + +def OLDnerd_tree(startpath, prec_seps=[]): + if not startpath.endswith('/'): + tosplit = startpath + startpath = startpath + '/' + else: + tosplit = startpath[:-1] + + rootnode = tosplit.split('/')[-1] + + items = nerd_tree_sort(startpath, dirfirst=True) + + ## rootnode is always a dir -> colorize + nerd_tree_txt = 'd_' + rootnode + + for n, item in enumerate(items): + if item in ITEMS_TO_EXCLUDE: + continue + + islast = (n==len(items) -1) + + sep = CHILD_CONNECTOR + if islast: + sep = LAST_CHILD_CONNECTOR + + if not islast: + psep_char = VERTICAL_CONNECTOR + else: + psep_char = '' + + if os.path.isdir(startpath+item): + seps = prec_seps[:] + seps.append((psep_char, '')) + treeline = produce_treeline(prec_seps, (sep, ' ')) + ' ' + nerd_tree(startpath+item, prec_seps=seps) + else: + treeline = produce_treeline(prec_seps, (sep, item)) + + nerd_tree_txt += '\n' + treeline + + return nerd_tree_txt + +def nerd_tree_from_struct(rootnode, prec_seps=[]): + ''' + rootnode is a dict + ''' + # items = nerd_tree_sort(startpath, dirfirst=True) + + ## rootnode is always a dir -> colorize + rootnode_name = rootnode['name'] + nerd_tree_txt = 'd_' + rootnode_name + + items = rootnode.get('children') or [] + + for n, item in enumerate(items): + if item['name'] in ITEMS_TO_EXCLUDE: + continue + + islast = (n==len(items) -1) + + sep = CHILD_CONNECTOR + if islast: + sep = LAST_CHILD_CONNECTOR + + if not islast: + psep_char = VERTICAL_CONNECTOR + else: + psep_char = '' + + if item['type'] == 'd': + seps = prec_seps[:] + seps.append((psep_char, '')) + treeline = produce_treeline(prec_seps, (sep, ' ')) + ' ' + nerd_tree_from_struct(item, prec_seps=seps) + else: + treeline = produce_treeline(prec_seps, (sep, item['name'])) + + nerd_tree_txt += '\n' + treeline + + return nerd_tree_txt + +def nerd_tree_struct(startpath, path=[]): + ''' + generate a recursive structure representing the tree + starting from startpath + ''' + if not startpath.endswith('/'): + tosplit = startpath + startpath = startpath + '/' + else: + tosplit = startpath[:-1] + + ## it is always a folder + rootnode = tosplit.split('/')[-1] + + # path_copy = path[:] + # path_copy.append(rootnode) + if path: + path.append({ + 'name' : rootnode, + 'type' : 'd', + 'abs_path' : startpath, + }) + else: + path = [{ + 'name' : './', + 'type' : 'd', + 'abs_path' : startpath, + }] + + d = { + 'name' : rootnode, + 'type' :'d', + 'abs_path' : startpath, + 'children' : [], + 'path': path, + } + + # using listdir + items = nerd_tree_sort(startpath, dirfirst=True) + + for item in items: + if item in ITEMS_TO_EXCLUDE: + continue + + if os.path.isdir(startpath+item): + d['children'].append(nerd_tree_struct(startpath+item, path=path[:])) + else: + path_copy = path[:] + path_copy.append({ + 'name' : item, + 'type' : 'f', + 'abs_path' : startpath + item + }) + + d['children'].append({ + 'name' : item, + 'type' : 'f', + 'abs_path' : startpath + item, + 'path' : path_copy, + }) + + return d + +def find_subtree(node, node_name, results=[]): + + if node['name'] == node_name: + results.append(node) + + children = node.get('children') or [] + + for item_child in children: + if item_child['type'] == 'd': + find_subtree(item_child, node_name, results) + else: + if item_child['name'] == node_name: + results.append(item_child) + +def find_node(startpath, item_name): + ''' + search item_name in subtree starting from startpath + Note: it finds all occurrencies + ''' + rnode = nerd_tree_struct(startpath) + + results = [] + find_subtree(rnode, item_name, results) + + return results + +def find_children(tree, item_abs_path): + ''' + It returns a tuple, first value is if father of node with item_name is found, + second value represents the children of parent node + ''' + if tree['abs_path'] == item_abs_path: + if not tree.get('children'): + tree['children'] = [] + return (True, tree['children']) + + children = tree.get('children') or [] + if not children: + return (False, []) + + for child in children: + found, fchildren = find_children(child, item_abs_path) + if found: + return (True, fchildren) + + return (False, []) + +def find_tree_struct(startpath, item_name): + ''' + Bad!!! + ''' + results = find_node(startpath, item_name) + if not results: + return {} + + tree_struct = {} + + for node in results: + path = node['path'] + + for n,p in enumerate(path): + pcopy = p.copy() + pname = p['name'] + + if pname == './': + if tree_struct: continue + tree_struct = pcopy + tree_struct['children'] = [] + continue + + parent = path[n-1] + found, children = find_children(tree_struct, parent['abs_path']) + + if not found: + raise Exception('Bug!!!') + + islastpath = True if (n== len(path)-1) else False + + if not islastpath: children.append(pcopy) + else : children.append(node) + + return tree_struct + +def nerd_tree(startpath): + tree_struct = nerd_tree_struct(startpath) + tree_txt = nerd_tree_from_struct(tree_struct) + print(tree_txt) + +def find_nerd_tree(startpath, item_name): + tree_struct = find_tree_struct(startpath, item_name) + if not tree_struct: + print('No results') + return + tree_txt = nerd_tree_from_struct(tree_struct) + print(tree_txt)