1
# Pytl
2
# Command line tool that parses python file, enumerates classes and methods as a tree structure.
3

4
import ast
5
import argparse
6

7
class LineText(object):
8
    def __init__(self, node, indent):
9
        self.node = node
10
        self.name = node.name
11
        self.indent = indent
12

13
    def get_line_number(self):
14
        if hasattr(self.node, 'lineno'):
15
            return self.node.lineno
16

17
    def get_indent_text(self):
18
        indent_text = ' ' * self.indent
19
        if self.indent != 0:
20
            indent_text += '|_'
21
        return indent_text
22

23

24
class ClassDefLineText(LineText):
25
    def __init__(self, class_node, indent):
26
        if not isinstance(class_node, ast.ClassDef):
27
            raise ValueError('node type is not ClassDef.')
28
        super(ClassDefLineText, self).__init__(class_node, indent)
29

30
    def __str__(self):
31
        return '{} : {}class {}'.format(self.get_line_number(), self.get_indent_text(), self.name)
32

33

34
class FunctionDefLineText(LineText):
35
    def __init__(self, def_node, indent):
36
        if not isinstance(def_node, ast.FunctionDef):
37
            raise ValueError('node type is not FunctionDef.')
38
        super(FunctionDefLineText, self).__init__(def_node, indent)
39
        try:
40
            self.args = ', '.join([a.arg for a in def_node.args.args])
41
        except AttributeError:
42
            # Support to python2
43
            self.args = ', '.join([a.id for a in def_node.args.args])
44

45
    def __str__(self):
46
        return '{} : {}def {}({})'.format(self.get_line_number(), self.get_indent_text(), self.name, self.args)
47

48
INDENT_NUM = 2
49

50
class Parser:
51

52
    def __init__(self, file_text, class_only=False, def_only=False):
53
        self.text_list = []
54
        self.class_only = class_only
55
        self.def_only = def_only
56
        self.tree = ast.parse(file_text)
57
        self.walk(self.tree)
58

59
    def walk(self, node, indent=0):
60
        next_indent = indent + INDENT_NUM
61
        if isinstance(node, ast.ClassDef) and not self.def_only:
62
            self.text_list.append(ClassDefLineText(node, indent))
63
        elif isinstance(node, ast.FunctionDef) and not self.class_only:
64
            self.text_list.append(FunctionDefLineText(node, indent))
65
        else:
66
            next_indent = indent
67

68
        for child in ast.iter_child_nodes(node):
69
            self.walk(child, indent=next_indent)
70

71
    def max_line_number_length(self):
72
        return len(str(max([x.get_line_number() for x in self.text_list])))
73

74
    def text_print(self):
75
        for line in self.text_list:
76
            left_space = self.max_line_number_length() - len(str(line.get_line_number()))
77
            print('{}{}'.format(' ' * left_space, line))
78

79

80
def main():
81
    parser = argparse.ArgumentParser(add_help=True)
82
    parser.add_argument('python_file_path', type=str, help='Specify target python file path')
83
    parser.add_argument('-c', '--class-only', action='store_true', default=False, help='Show only the class line')
84
    parser.add_argument('-d', '--def-only', action='store_true', default=False, help='Show only the def line')
85
    parse_args = parser.parse_args()
86
    with open(parse_args.python_file_path, 'r') as f:
87
        Parser(f.read(), parse_args.class_only, parse_args.def_only).text_print()
88

89
if __name__ == '__main__':
90
    main()