tree, tree_find refactoring

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))
+

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))
+