import time import re import fnmatch from src.colors import RED, YELLOW, GREEN, CYAN , BLUE, PURPLE from src.tree import NerdTree class NerdTreeFind(NerdTree): ''' @param startpath string: the startpath from which to generate the json_tree member (the tree representation) @param opts dict: the opts for the tree representation ''' def __init__(self, startpath, opts={}, find_opts={}): self.item_name = '' self.re_item_name = None ''' re_item_name is a regexp python object, can be 1. a regular expression, specifying in the command line -re option 2. a string that supports the shell expansion ''' self.results = [] self.start_time = time.time() self.end_time = None ## computed tree for find self.json_tree_find = {} self.find_opts = find_opts NerdTree.__init__(self, startpath, opts=opts) def found_true_cond(self, test_node): search_type = self.find_opts['type'] if self.find_opts.get('type') else None if not search_type: test_type = True else: test_type = (test_node['type'] == search_type) test_match = True if self.re_item_name.match(test_node['name']) else False return all([test_match, test_type]) def find_node_recursively(self, node): if self.found_true_cond(node): node_copy = node.copy() node_copy.update({'color_formatter' : YELLOW, 'found' : 1}) self.results.append(node_copy) children = node.get('children') or [] for item_child in children: if item_child['type'] == 'd': self.find_node_recursively(item_child) else: if self.found_true_cond(item_child): item_child_copy = item_child.copy() item_child_copy.update({'color_formatter' : YELLOW, 'found' : 1}) self.results.append(item_child_copy) def find_node(self): self.find_node_recursively(self.json_tree) return self.results def find_children(self, tree, item_abs_path): ''' It returns a tuple: first value True if father of node with item_abs_path is found, second value represents the children of parent node with item_abs_path ''' 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 = self.find_children(child, item_abs_path) if found: return (True, fchildren) return (False, []) def find_tree_struct(self): ''' recreates the tree with search results ''' results = self.find_node() ## remove from this level the option if not results: return {} ## filter results items_to_include= self.find_opts['items_to_include'] if self.find_opts.get('items_to_include') is not None else [] if items_to_include: filtered_results = [] items_to_include_set = set(items_to_include) for res in results: pathset = set(path['name'] for path in res['path']) if not items_to_include_set.intersection(pathset): continue filtered_results.append(res.copy()) results = filtered_results dont_show_children_nodes = self.find_opts['dont_show_children_nodes'] if self.find_opts.get('dont_show_children_nodes') is not None else False show_subtree_info = self.find_opts['show_subtree_info'] if self.find_opts.get('show_subtree_info') is not None else False # self.opts['dont_show_children_nodes'] = dont_show_children_nodes # self.opts['show_subtree_info'] = show_subtree_info 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 = self.find_children(tree_struct, parent['abs_path']) if not found: raise Exception('Bug!!!') if p['abs_path'] in [c['abs_path'] for c in children]: continue islastpath = True if (n== len(path)-1) else False if not islastpath: children.append(pcopy) else : children.append(node) self.compute_aggregate_recursively(tree_struct) return tree_struct def get_reobj_find(self, item_name): if not self.find_opts['use_regular_expression']: regex = fnmatch.translate(item_name) reobj = re.compile(regex) return reobj reobj = re.compile(r'^%s$' % (item_name,)) return reobj def reset(self, new_item_name, new_opts): self.find_opts = new_opts self.results = [] self.item_name = new_item_name print('Starting from path: %s' % CYAN(self.startpath,)) print('Use regexp: %s' % (CYAN(self.find_opts['use_regular_expression']))) print('pattern: %s' % CYAN(new_item_name,)) try: self.re_item_name = self.get_reobj_find(new_item_name) except re.error: raise Exception('Can\'t compile regexp %s, exit' % (new_item_name,)) self.json_tree_find = {} def print_elapsed_time(self): self.end_time = time.time() self.elapsed_time = round(self.end_time - self.start_time,2) print('\nElapsed Time: %s' % CYAN(self.elapsed_time)) def print(self, item_name='', opts={}): if not item_name and not self.item_name: print('Please use a valid search string') return if item_name != self.item_name or self.find_opts != opts: self.reset(item_name, opts) self.json_tree_find = self.find_tree_struct() if not self.json_tree_find: print('No results') self.print_elapsed_time() return tree_res = self.tree_from_struct(self.json_tree_find) print(tree_res) self.print_elapsed_time() def find(self, item_name, opts={}): ''' @param item_name the string to search in the tree @param opts dict for the find option ''' return self.print(item_name, opts=opts)