1#!/usr/bin/python3 2 3import os 4 5class TextTree(object): 6 7 def tostring(self, object): 8 strings = self.tostrings(object) 9 return ''.join(strings) 10 11 def tostrings(self, object): 12 strings = [object.name + '\n'] 13 14 def add_item(item, ps1, ps2): 15 if isinstance(item, str): 16 strings.append(ps1 + item + '\n') 17 else: 18 substrings = self.tostrings(item) # recurs 19 strings.append(ps1 + substrings.pop(0)) 20 strings.extend([ps2 + s for s in substrings]) 21 22 items = list(object.items()) 23 if items: 24 for i in range(len(items) - 1): 25 add_item(items[i], '|-- ', '| ') 26 add_item(items[-1], '`-- ', ' ') 27 28 return strings 29 30 31 32class ModuleFile(object): 33 34 def __init__(self, file): 35 assert os.path.isfile(file), 'Could not find file: %s' % file 36 self.file = file 37 self.name = os.path.basename(file)[:-3] 38 39 self.classes = [] 40 for line in open(self.file): 41 line = line.strip() 42 if line.startswith('class') and line.endswith(':'): 43 self.classes.append(line[5:-1].strip()) 44 45 def items(self): 46 return self.classes[:] 47 48 49class ModuleDir(ModuleFile): 50 51 def __init__(self, dir): 52 assert os.path.isdir(dir), 'Could not find dir: %s' % dir 53 ModuleFile.__init__(self, dir + '/__init__.py') 54 self.dir = dir 55 self.name = os.path.basename(dir) 56 self.modules = [] 57 58 paths = [dir + '/' + p for p in os.listdir(dir) if not p.startswith('_')] 59 for file in [f for f in paths if f.endswith('.py')]: 60 self.modules.append(ModuleFile(file)) 61 for subdir in [d for d in paths if os.path.isdir(d)]: 62 self.modules.append(ModuleDir(subdir)) 63 64 self.modules.sort(key=lambda m: m.name) 65 66 def items(self): 67 items = ModuleFile.items(self) 68 items.extend(self.modules) 69 return items 70 71 72if __name__ == '__main__': 73 dir = ModuleDir('./zim') 74 print(TextTree().tostring(dir)) 75