tree, tree_find refactoring

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