ntree colors

ntree.py

@@ -9,11 +9,13 @@ if __name__ == '__main__':
   # print(cmd_args.exclude)
   path = cmd_args.path
   list_dir_first = cmd_args.list_directories_first
+  colorize = cmd_args.colorize
 
   # opts for generatign the tree
   opts = {
       'items_to_exclude' : cmd_args.exclude or [],
-      'list_dir_first' : list_dir_first
+      'list_dir_first' : list_dir_first,
+      'colorize' : colorize,
       }
 
   ntree = NerdTree(path, opts=opts)

src/ntree_cmdline_parser.py

@@ -12,6 +12,9 @@ cmd_parser.add_argument("path")
 cmd_parser.add_argument("-e", "--exclude", action="append")
 cmd_parser.add_argument("-l", "--list-directories-first", action="store_true")
 
+# colorize
+cmd_parser.add_argument("-c", "--colorize", action="store_true")
+
 # start the command line parsing
 cmd_args = cmd_parser.parse_args()
 

src/tree.py

@@ -1,4 +1,5 @@
 import os
+import stat as _stat
 from pathlib import Path
 
 from src.colors import RED, YELLOW, GREEN, CYAN , BLUE, PURPLE
@@ -64,10 +65,24 @@ class NerdTree():
 
     return os.listdir(startpath)
 
+  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
+    t = d['type']
+    d['color_formatter'] = _tree_repr.get_color_formatter(t)
+    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
@@ -103,7 +118,7 @@ class NerdTree():
         'is_symlink' : start_dir_is_symlink,
         }]
 
-    d = {
+    d = self.colorize({
       'name' : rootnode,
       'type' :'d',
       'abs_path' : startpath,
@@ -112,7 +127,7 @@ class NerdTree():
       'size' : stat_dir_item.st_size,
       'target' : str(start_dir_target),
       'is_symlink' : start_dir_is_symlink,
-      }
+      })
 
     # using listdir
     try:
@@ -129,32 +144,35 @@ class NerdTree():
       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()
+
+        stat_item = path_object.lstat()
+        stat_item_mode = stat_item.st_mode
+
         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,
-          })
+        is_executable = stat_item_mode & _stat.S_IXUSR
 
-        d['children'].append({
+        path_copy = path[:]
+
+        it_dict = {
           '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,
-          })
+          '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
 

src/tree_repr.py

@@ -1,6 +1,6 @@
 import os
 
-# from colors import RED, YELLOW, GREEN, CYAN , BLUE, PURPLE
+from src.colors import RED, YELLOW, GREEN, CYAN , BLUE, PURPLE
 
 EMPTY_TREE_LINE_LENGTH = 160
 EMPTY_TREE_LINE = ' ' * EMPTY_TREE_LINE_LENGTH
@@ -71,3 +71,15 @@ def format(num, unit=None):
   # human_number = num
   
   return human_number, unit
+
+
+def get_color_formatter(t):
+  '''
+    @param        t      string can be 'd', 'f', 'l', 'x'
+  '''
+  if t == 'd':
+    return BLUE
+  if t == 'x':
+    return GREEN
+  return None
+