338 lines
9.4 KiB
Python
338 lines
9.4 KiB
Python
import os
|
|
import stat as _stat
|
|
import fnmatch, re
|
|
|
|
from pathlib import Path
|
|
|
|
from src.colors import RED, YELLOW, GREEN, CYAN , BLUE, PURPLE
|
|
import src.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
|
|
'''
|
|
the below member contains regexp compiled for the items to exclude
|
|
the match in any of these regexp must cause the skip
|
|
'''
|
|
self.reobj_excludes = []
|
|
self.init_reobj_excludes()
|
|
if startpath:
|
|
self.json_tree = self.tree_struct(startpath)
|
|
else:
|
|
self.json_tree = {}
|
|
|
|
def init_reobj_excludes(self):
|
|
'''
|
|
it use bash expansion pattern to generate
|
|
regexp to check for exclude items
|
|
'''
|
|
exclusions = self.opts.get('items_to_exclude') or []
|
|
if not exclusions:
|
|
return
|
|
|
|
for e in exclusions:
|
|
# for example e -> '*.txt'
|
|
regex = fnmatch.translate(e)
|
|
reobj = re.compile(regex)
|
|
self.reobj_excludes.append(reobj)
|
|
|
|
|
|
def get_path_obj(self, abs_path_item):
|
|
## using lstat not produce error in case of symbolic link
|
|
path_object = Path(abs_path_item)
|
|
|
|
item_type = ''
|
|
|
|
if path_object.is_dir():
|
|
item_type = 'd'
|
|
elif path_object.is_symlink():
|
|
item_type = 'l'
|
|
elif path_object.is_file():
|
|
item_type = 'f'
|
|
|
|
assert item_type
|
|
|
|
return (item_type, path_object)
|
|
|
|
def sort_directories_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_directories_first(startpath)
|
|
|
|
return os.listdir(startpath)
|
|
|
|
def test_exclude_true(self, item_name):
|
|
'''
|
|
@param item_name string - the name of file or folder
|
|
'''
|
|
if not self.reobj_excludes:
|
|
return False
|
|
|
|
reobj_tests = list(map(lambda reobj: reobj.match(item_name) ,self.reobj_excludes))
|
|
if any(reobj_tests):
|
|
return True
|
|
|
|
return False
|
|
|
|
def colorize(self, d):
|
|
'''
|
|
use the mandatory key `type` to format item
|
|
'''
|
|
colorize = self.opts['colorize'] if self.opts.get('colorize') is not None else False
|
|
if not colorize:
|
|
return d
|
|
|
|
if d.get('is_symlink'):
|
|
d['color_formatter'] = _tree_repr.get_color_formatter('l')
|
|
d['target_color_formatter'] = _tree_repr.get_color_formatter(d['type'])
|
|
else:
|
|
d['color_formatter'] = _tree_repr.get_color_formatter(d['type'])
|
|
d['target_color_formatter'] = ''
|
|
|
|
return d
|
|
|
|
def tree_struct(self, startpath, path=[]):
|
|
'''
|
|
generate a recursive structure representing the tree
|
|
starting from startpath
|
|
|
|
NOTE: here is the correct place to put metadata useful
|
|
for formatting the output
|
|
'''
|
|
if not startpath.endswith('/'):
|
|
tosplit = startpath
|
|
startpath = startpath + '/'
|
|
else:
|
|
tosplit = startpath[:-1]
|
|
|
|
## rootnode for definition is always a folder
|
|
rootnode = tosplit.split('/')[-1]
|
|
|
|
start_dir_item_type, start_dir_object = self.get_path_obj(startpath)
|
|
stat_dir_item = start_dir_object.lstat()
|
|
|
|
start_dir_is_symlink = start_dir_object.is_symlink()
|
|
start_dir_target = start_dir_object.resolve() if start_dir_is_symlink else ''
|
|
|
|
if path:
|
|
path.append({
|
|
'name' : rootnode,
|
|
'type' : start_dir_item_type,
|
|
'abs_path' : startpath,
|
|
'size' : stat_dir_item.st_size,
|
|
'target' : str(start_dir_target),
|
|
'is_symlink' : start_dir_is_symlink,
|
|
})
|
|
else:
|
|
path = [{
|
|
'name' : './',
|
|
'type' : start_dir_item_type,
|
|
'abs_path' : startpath,
|
|
'size' : stat_dir_item.st_size,
|
|
'target' : str(start_dir_target),
|
|
'is_symlink' : start_dir_is_symlink,
|
|
}]
|
|
|
|
d = self.colorize({
|
|
'name' : rootnode,
|
|
'type' :'d',
|
|
'abs_path' : startpath,
|
|
'children' : [],
|
|
'path': path,
|
|
'size' : stat_dir_item.st_size,
|
|
'target' : str(start_dir_target),
|
|
'is_symlink' : start_dir_is_symlink,
|
|
})
|
|
|
|
follow_symbolic_link = self.opts['follow_symbolic_link'] if self.opts.get('follow_symbolic_link') is not None else False
|
|
if start_dir_is_symlink and not follow_symbolic_link:
|
|
return d
|
|
|
|
# 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 (self.opts.get('items_to_exclude') or []):
|
|
# continue
|
|
|
|
if self.test_exclude_true(item):
|
|
continue
|
|
|
|
|
|
abs_path_item = startpath+item
|
|
item_type, path_object = self.get_path_obj(abs_path_item)
|
|
|
|
if path_object.is_dir():
|
|
d['children'].append(self.tree_struct(abs_path_item, path=path[:]))
|
|
else:
|
|
is_symlink = path_object.is_symlink()
|
|
stat_item = path_object.lstat()
|
|
stat_item_mode = stat_item.st_mode
|
|
|
|
target = path_object.resolve() if is_symlink else ''
|
|
|
|
is_executable = stat_item_mode & _stat.S_IXUSR
|
|
|
|
path_copy = path[:]
|
|
|
|
it_dict = {
|
|
'name' : item,
|
|
'type' : item_type,
|
|
'abs_path' : abs_path_item,
|
|
'target' : str(target),
|
|
'is_symlink' : is_symlink,
|
|
'type' : 'x' if is_executable else None,
|
|
}
|
|
|
|
path_copy.append(it_dict)
|
|
|
|
enhanced = dict(it_dict, path=path_copy, size = stat_item.st_size)
|
|
|
|
d['children'].append(self.colorize(enhanced))
|
|
|
|
return d
|
|
|
|
def compute_aggregate_recursively(self, node=None):
|
|
'''
|
|
node is the rootnode of the subtree
|
|
at the moment compute size and number of items of subtree starting from node
|
|
'''
|
|
|
|
subtree = {
|
|
'subtree_size' : 0,
|
|
'subtree_items' : 0,
|
|
}
|
|
|
|
node = node or self.json_tree
|
|
|
|
for item in node.get('children') or []:
|
|
|
|
if item['type'] == 'd':
|
|
if item.get('subtree_size') is None:
|
|
subtree_compute = self.compute_aggregate_recursively(item)
|
|
else:
|
|
subtree_compute = {
|
|
'subtree_size' : item['subtree_size'],
|
|
'subtree_items' : item['subtree_items'],
|
|
}
|
|
## if you want count the current directory as item you must sum 1
|
|
## if you want count the current directory size you must include item['size']
|
|
subtree['subtree_size'] += subtree_compute['subtree_size'] + item['size']
|
|
subtree['subtree_items'] += subtree_compute['subtree_items'] + 1
|
|
else:
|
|
subtree['subtree_size'] += item['size']
|
|
subtree['subtree_items'] += 1
|
|
|
|
# return subtree
|
|
node.update(subtree)
|
|
return subtree
|
|
|
|
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
|
|
|
|
if (rootnode['is_symlink'] and rootnode.get('target', '')):
|
|
if rootnode.get('target_color_formatter'):
|
|
nerd_tree_txt += ' -> %s' % (rootnode['target_color_formatter'](rootnode['target']),)
|
|
else:
|
|
nerd_tree_txt += ' -> %s' % (rootnode['target'],)
|
|
|
|
items = rootnode.get('children') or []
|
|
|
|
nodefound = rootnode.get('found')
|
|
if nodefound:
|
|
if self.find_opts['dont_show_children_nodes'] or self.find_opts['show_subtree_info']:
|
|
if rootnode.get('subtree_items'):
|
|
hnum, unit = _tree_repr.format(rootnode['subtree_size'])
|
|
nerd_tree_txt += ' (%s item(s)/%s%s ▾)' % (rootnode['subtree_items'], hnum, unit)
|
|
|
|
if self.find_opts['dont_show_children_nodes']: items = []
|
|
|
|
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))
|
|
if item.get('found') and self.find_opts['show_subtree_info']:
|
|
hnum, unit = _tree_repr.format(item['size'])
|
|
treeline += ' (%s%s)' % (hnum, unit)
|
|
|
|
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))
|
|
|