diff --git a/src/__pycache__/tree.cpython-39.pyc b/src/__pycache__/tree.cpython-39.pyc new file mode 100644 index 0000000..eb9749b Binary files /dev/null and b/src/__pycache__/tree.cpython-39.pyc differ diff --git a/src/__pycache__/tree_find.cpython-39.pyc b/src/__pycache__/tree_find.cpython-39.pyc new file mode 100644 index 0000000..6e3c145 Binary files /dev/null and b/src/__pycache__/tree_find.cpython-39.pyc differ diff --git a/src/tree.py b/src/tree.py new file mode 100644 index 0000000..b0aefc3 --- /dev/null +++ b/src/tree.py @@ -0,0 +1,171 @@ +import os +from colors import RED, YELLOW, GREEN, CYAN , BLUE, PURPLE + +import tree_repr as _tree_repr + +class NerdTree(): + ''' + @param startpath string + @param opts dict of options + + list_dir_first + ''' + + def __init__(self, startpath='', opts={}): + self.startpath = startpath + self.opts = opts + if startpath: + self.json_tree = self.tree_struct(startpath) + else: + self.json_tree = {} + + def sort_directory_first(self, 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 sort_items(self, startpath): + ''' + it returns the sorted list of items starting from startpath + ''' + + # dirs = [] + # files = [] + + dirfirst = self.opts['list_dir_first'] if self.opts.get('list_dir_first') is not None else False + + if dirfirst: + return self.sort_directory_first(startpath) + + return os.listdir(startpath) + + def tree_struct(self, startpath, path=[]): + ''' + generate a recursive structure representing the tree + starting from startpath + ''' + if not startpath.endswith('/'): + tosplit = startpath + startpath = startpath + '/' + else: + tosplit = startpath[:-1] + + ## rootnode for definition is always a folder + rootnode = tosplit.split('/')[-1] + + 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 + try: + items = self.sort_items(startpath) + except Exception as ss: + print(f'Path: {startpath} not readable because {ss}') + items = [] + d.setdefault('not_readable', 1) + + for item in items: + if item in _tree_repr.ITEMS_TO_EXCLUDE: + continue + + if os.path.isdir(startpath+item): + d['children'].append(self.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 tree_from_struct(self, rootnode, prec_seps=[]): + ''' + recursively produce a string for representing the tree + rootnode is a node dict + ''' + + ## rootnode is always a dir -> colorize + rootnode_name = rootnode['name'] + if rootnode.get('color_formatter'): + nerd_tree_txt = rootnode['color_formatter'](rootnode_name) + else: + nerd_tree_txt = rootnode_name + + items = rootnode.get('children') or [] + + for n, item in enumerate(items): + if item['name'] in _tree_repr.ITEMS_TO_EXCLUDE: + continue + + islast = (n==len(items) -1) + + sep = _tree_repr.CHILD_CONNECTOR + if islast: + sep = _tree_repr.LAST_CHILD_CONNECTOR + + if not islast: + psep_char = _tree_repr.VERTICAL_CONNECTOR + else: + psep_char = '' + + if item['type'] == 'd': + seps = prec_seps[:] + seps.append((psep_char, {'name' : ''})) + treeline = _tree_repr.produce_treeline(prec_seps, (sep, {'name' : ' '})) + ' ' + self.tree_from_struct(item, prec_seps=seps) + else: + treeline = _tree_repr.produce_treeline(prec_seps, (sep, item)) + + nerd_tree_txt += '\n' + treeline + + return nerd_tree_txt + + def print(self, startpath='', opts={}): + reinit = False + + if opts and self.opts != opts: + self.opts = opts + reinit = True + + if startpath and startpath != self.startpath: + self.startpath = startpath + reinit = True + + if reinit: + self.json_tree = self.tree_struct(self.startpath) + + print(self.tree_from_struct(self.json_tree)) + diff --git a/src/tree_find.py b/src/tree_find.py new file mode 100644 index 0000000..b82ff9e --- /dev/null +++ b/src/tree_find.py @@ -0,0 +1,104 @@ +from colors import RED, YELLOW, GREEN, CYAN , BLUE, PURPLE +from tree import NerdTree + +class NerdTreeFind(NerdTree): + + def __init__(self, startpath, item_name, opts={}): + self.item_name = item_name + self.results = [] + NerdTree.__init__(self, startpath, opts=opts) + + def find_node_recursively(self, node): + if node['name'] == self.item_name: + node.update({'color_formatter' : YELLOW}) + self.results.append(node) + + children = node.get('children') or [] + + for item_child in children: + if item_child['type'] == 'd': + self.find_node_recursively(item_child) + else: + if item_child['name'] == self.item_name: + item_child.update({'color_formatter' : YELLOW}) + self.results.append(item_child) + + def find_node(self): + self.find_node_recursively(self.json_tree) + return self.results + + def find_children(self, tree, item_abs_path): + ''' + It returns a tuple: + first value True if father of node with item_abs_path is found, + second value represents the children of parent node with item_abs_path + ''' + 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 = self.find_children(child, item_abs_path) + if found: + return (True, fchildren) + + return (False, []) + + def find_tree_struct(self, item_name): + ''' + ''' + results = self.find_node() + + dont_show_children_nodes = self.opts['dont_show_children_nodes'] if self.opts.get('dont_show_children_nodes') else False + + if not results: + return {} + + if dont_show_children_nodes: + for result in results: + result['children'] = [] + + 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 = self.find_children(tree_struct, parent['abs_path']) + + if not found: + raise Exception('Bug!!!') + + if p['abs_path'] in [c['abs_path'] for c in children]: + continue + + islastpath = True if (n== len(path)-1) else False + + if not islastpath: children.append(pcopy) + else : children.append(node) + + return tree_struct + + def print(self): + self.tree_struct = tree_struct = self.find_tree_struct(self.item_name) + if not tree_struct: + print('No results') + return + + print(self.tree_from_struct(self.tree_struct)) +