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