import os from pathlib import Path from src.colors import RED, YELLOW, GREEN, CYAN , BLUE, PURPLE import src.tree_repr as _tree_repr import src.tree_find as _tree_find class NerdTree(): ''' @param startpath string @param opts dict of options list_dir_first ''' def __init__(self, startpath='', opts={}): self.startpath = startpath self.opts = opts if startpath: self.json_tree = self.tree_struct(startpath) else: self.json_tree = {} 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 tree_struct(self, startpath, path=[]): ''' generate a recursive structure representing the tree starting from startpath ''' 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 = { '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, } # 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 abs_path_item = startpath+item item_type, path_object = self.get_path_obj(abs_path_item) stat_item = path_object.lstat() if path_object.is_dir(): d['children'].append(self.tree_struct(abs_path_item, path=path[:])) else: is_symlink = path_object.is_symlink() target = path_object.resolve() if is_symlink else '' path_copy = path[:] path_copy.append({ 'name' : item, 'type' : item_type, 'abs_path' : abs_path_item, 'target' : str(target), 'is_symlink' : is_symlink, }) d['children'].append({ 'name' : item, 'type' : item_type, 'abs_path' : abs_path_item, 'path' : path_copy, 'size' : stat_item.st_size, 'target' : str(target), 'is_symlink' : is_symlink, }) 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', '')): nerd_tree_txt += ' -> %s' % (rootnode['target'],) items = rootnode.get('children') or [] if isinstance(self, _tree_find.NerdTreeFind): 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 isinstance(self, _tree_find.NerdTreeFind): 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))