tree_repr base

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)