tree_repr base
This commit is contained in:
Amber
parent
e5c7abe436
commit
665c264315
|
|
@ -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)
|
||||
Loading…
Reference in New Issue