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