PythonLanguageStrategy.get_imports()   A
last analyzed

Complexity

Conditions 3

Size

Total Lines 32
Code Lines 22

Duplication

Lines 0
Ratio 0 %

Importance

Changes 0
Metric Value
cc 3
eloc 22
nop 2
dl 0
loc 32
rs 9.352
c 0
b 0
f 0
1
from abc import ABC, abstractmethod
2
3
from antlr4 import InputStream  # type: ignore
4
from antlr4 import CommonTokenStream, ParserRuleContext, ParseTreeWalker
5
from antlr4.Token import CommonToken  # type: ignore
6
7
from menderbot.antlr_generated.PythonLexer import PythonLexer  # type: ignore
8
from menderbot.antlr_generated.PythonParser import PythonParser  # type: ignore
9
from menderbot.antlr_generated.PythonParserListener import (  # type: ignore
10
    PythonParserListener,
11
)
12
13
14
def node_str(node) -> str:
15
    return get_text_including_whitespace(node)
16
17
18
def node_start_line(node: ParserRuleContext) -> int:
19
    node_start: CommonToken = node.start
20
    return node_start.line
21
22
23
def node_stop_line(node: ParserRuleContext) -> int:
24
    node_stop: CommonToken = node.stop
25
    return node_stop.line
26
27
28
def node_start_column(node: ParserRuleContext) -> int:
29
    node_start: CommonToken = node.start
30
    return node_start.column
31
32
33
def node_stop_column(node: ParserRuleContext) -> int:
34
    node_start: CommonToken = node.stop
35
    return node_start.column
36
37
38
def get_text_including_whitespace(ctx: ParserRuleContext):
39
    if not ctx:
40
        return ""
41
    start: CommonToken = ctx.start
42
    stop: CommonToken = ctx.stop
43
    start_idx: int = start.start
44
    stop_idx: int = stop.stop
45
    stream: InputStream = start.getInputStream()
46
    return stream.getText(start_idx, stop_idx)
47
48
49
def line_indent(line: str) -> str:
50
    count = len(line) - len(line.lstrip())
51
    return line[:count]
52
53
54
def function_indent(code: str) -> str:
55
    second_line_start = code.find("\n") + 1
56
    no_first_line = code[second_line_start:]
57
    if no_first_line.find("\n") > -1:
58
        second_line_end = second_line_start + no_first_line.find("\n")
59
        second_line = code[second_line_start:second_line_end]
60
    else:
61
        second_line = no_first_line
62
    return line_indent(second_line)
63
64
65
def reindent(text: str, indent: str) -> str:
66
    lines = text.split("\n")
67
    indented_lines = [indent + line.lstrip() for line in lines]
68
    return "\n".join(indented_lines)
69
70
71
class LanguageStrategy(ABC):
72
    @abstractmethod
73
    def function_has_comment(self, node) -> bool:
74
        pass
75
76
    @abstractmethod
77
    def parse_source_to_tree(self, source: bytes) -> None:
78
        pass
79
80
    @abstractmethod
81
    def get_function_nodes(self, tree) -> list:
82
        pass
83
84
    def get_imports(self, tree) -> list:
85
        del tree
86
        return []
87
88
    @property
89
    @abstractmethod
90
    def function_doc_line_offset(self) -> int:
91
        pass
92
93
94
class PythonLanguageStrategy(LanguageStrategy):
95
    def function_has_comment(self, node: PythonParser.FuncdefContext) -> bool:
96
        """Checks if function has a docstring."""
97
        body_node: PythonParser.SuiteContext = node.suite()
98
        if body_node:
99
            # https://peps.python.org/pep-0257/
100
            first_stmt_node: PythonParser.StmtContext = body_node.stmt(0)
101
            first_stmt_text = first_stmt_node.getText().strip()
102
            has_doc_prefix = (
103
                first_stmt_text.startswith('"""')
104
                or first_stmt_text.startswith('r"""')
105
                or first_stmt_text.startswith('u"""')
106
            )
107
108
            has_doc_suffix = first_stmt_text.endswith('"""')
109
            return has_doc_prefix and has_doc_suffix
110
        return False
111
112
    def parse_source_to_tree(self, source: bytes):
113
        input_stream = InputStream(str(source, encoding="utf-8") + "\n")
114
        lexer = PythonLexer(input_stream)
115
        token_stream = CommonTokenStream(lexer)
116
        parser = PythonParser(token_stream)
117
        return parser.file_input()
118
119
    def get_function_nodes(self, tree) -> list[PythonParser.FuncdefContext]:
120
        function_nodes: list[PythonParser.FuncdefContext] = []
121
122
        class MyListener(PythonParserListener):
123
            def enterFuncdef(self, ctx: PythonParser.FuncdefContext):
124
                function_nodes.append(ctx)
125
                # print(ctx.toStringTree(recog=parser))
126
127
        walker = ParseTreeWalker()
128
        walker.walk(MyListener(), tree)
129
        return function_nodes
130
131
    def get_function_node_name(self, node):
132
        name_node: PythonParser.NameContext = node.name()
133
        name = name_node.getText()
134
        return name
135
136
    def get_imports(self, tree) -> list[tuple[str, str]]:
137
        results: list[tuple[str, str]] = []
138
139
        class MyListener(PythonParserListener):
140
            def enterFrom_stmt(self, ctx: PythonParser.From_stmtContext):
141
                dotted_name_ctx: PythonParser.Dotted_nameContext = ctx.dotted_name()
142
                import_as_names_ctx: PythonParser.Import_as_namesContext = (
143
                    ctx.import_as_names()
144
                )
145
                import_as_name_ctxs: list[PythonParser.Import_as_nameContext] = (
146
                    import_as_names_ctx.import_as_name()
147
                )
148
                # print(type(import_as_name_ctxs))
149
                for import_as_name_ctx in import_as_name_ctxs:
150
                    results.append(
151
                        (dotted_name_ctx.getText(), node_str(import_as_name_ctx))
152
                    )
153
                # print(ctx.toStringTree(recog=parser))
154
155
            def enterImport_stmt(self, ctx: PythonParser.Import_stmtContext):
156
                dotted_as_names_ctx: PythonParser.Dotted_as_namesContext = (
157
                    ctx.dotted_as_names()
158
                )
159
                dotted_as_name_ctxs: list[PythonParser.Dotted_nameContext] = (
160
                    dotted_as_names_ctx.dotted_as_name()
161
                )
162
                for dottedAsNameCtx in dotted_as_name_ctxs:
163
                    results.append(("", node_str(dottedAsNameCtx)))
164
165
        walker = ParseTreeWalker()
166
        walker.walk(MyListener(), tree)
167
        return results
168
169
    function_doc_line_offset = 1
170
171
172
# class CppLanguageStrategy(LanguageStrategy):
173
#     def function_has_comment(self, node) -> bool:
174
#         return node.prev_sibling.type in ["comment"]
175
#
176
#     def parse_source_to_tree(self, source: bytes):
177
#         return parse_source_to_tree(source, CPP_LANGUAGE)
178
#
179
#     def get_function_nodes(self, tree) -> list:
180
#         query = CPP_LANGUAGE.query(
181
#             """
182
#         [
183
#             (function_definition) @function
184
#         ]
185
#         """
186
#         )
187
#         captures = query.captures(tree.root_node)
188
#         return [capture[0] for capture in captures]
189
#
190
#     function_doc_line_offset = 0
191
192
193
LANGUAGE_STRATEGIES = {
194
    ".py": PythonLanguageStrategy(),
195
    # ".c": CppLanguageStrategy(),
196
    # ".cpp": CppLanguageStrategy(),
197
}
198