Coverage for src / renaissance / integrations / python / ast / rst_node.py: 94%
324 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-09-09 14:04 +0000
1import ast
2import sys
3import textwrap
4from collections.abc import Callable, Sequence
5from pathlib import Path
6from typing import Any, Self
8from renaissance.integrations.python.ast.util import convert
9from renaissance.integrations.types import KIND_MAP, OPERATOR_MAP, Assert, FunctionDef, Global, ImplicitNode, Tuple, UnknownType
10from renaissance.syntax_tree.match_finder import find_in_list
11from renaissance.utils.ast_utils import (
12 format_node,
13 match_children,
14 match_props,
15 next_sibling,
16 preceding_sibling,
17 traverse,
18)
20types = ["int", "float", "str", "list", "set", "tuple", "Mapping", "dict", "Optional"]
21IRRELEVANT_PROPS = {"comment"}
22IRRELEVANT_NODES = {"comment"}
23IMPLICIT = [ImplicitNode]
26class ImplicitNode(ast.Name):
27 _fields = (
28 "id",
29 "body",
30 )
32 _field_types = {
33 "id": str,
34 "body": list,
35 }
37 def __init__(self, name, children=None):
38 super().__init__(name, children or [])
39 self.lineno = 0
40 self.col_offset = 0
41 self.end_lineno = 0
42 self.end_col_offset = 0
45class PythonRSTReference:
46 def __repr__(self):
47 return f"{self.node_id}:{self.ref_kind}"
49 def __init__(self, node_id: str, ref_kind: str, properties: dict[str, Any]) -> None:
50 self.node_id = node_id
51 self.ref_kind = ref_kind
52 self.properties = properties
55class PythonRstTranslationUnit:
56 cache = {}
58 def __init__(self, content, file_name: str):
59 self.content = content.encode(sys.getfilesystemencoding())
60 self.atu = ast.parse(content, file_name)
61 self.file_name = file_name
62 self.references_initialized = False
63 PythonRstTranslationUnit.cache[file_name] = content
64 self.lines = self.content.splitlines()
66 self._references: dict[str, list[PythonRSTReference]] = {}
67 self._referenced_by: dict[str, list[PythonRSTReference]] = {}
68 self._nodes: dict[str, PythonRstNode] = {}
70 def check_diagnostics(self, continue_with_warning=True) -> None:
71 msg = None
72 errors = ""
73 for d in self.atu.type_ignores:
74 msg = f"type ignored: {d.tag} at {d.lineno}\n"
75 errors += msg
76 print(msg)
77 if msg and not continue_with_warning:
78 raise Exception(f"Error parsing: {self.file_name} \n+ errors: {errors}")
80 def lazy_create_refers(self, node: PythonRstNode) -> None:
81 if self.references_initialized:
82 return
83 for n in traverse(node.root):
84 self.create_references(n)
85 self.references_initialized = True
87 def add(self, node):
88 match node.ast_type.__name__:
89 case "Name":
90 if node.node.id not in self._nodes and node.node.id not in types:
91 self._nodes[node.node.id] = node
92 case "FunctionDef":
93 if node.node.name not in self._nodes:
94 self._nodes[node.node.name] = node
95 case "Call":
96 if node.name not in self._nodes:
97 self._nodes[node.name] = node
98 case "ClassDef":
99 if node.name not in self._nodes:
100 self._nodes[node.name] = node
101 case "arg":
102 if node.name != "self" and node.name not in self._nodes:
103 self._nodes[node.name] = node
105 def create_references(self, ast_node) -> None:
106 assert isinstance(ast_node, PythonRstNode), f"Expected PythonASTNode but got {type(ast_node)}"
107 match type(ast_node.node):
108 case ast.arg:
109 if ast_node.name != "self" and isinstance(ast_node.node, ast.arg) and isinstance(ast_node.node.annotation, ast.Name):
110 node_id = ast_node.name
111 ref_id = ast_node.node.annotation.id
112 ref_kind = "TypeRef"
113 self.add_reference(node_id, ref_id, ref_kind)
114 case ast.Assign:
115 if isinstance(ast_node.node, ast.Assign):
116 for n in ast_node.node.targets:
117 if isinstance(n, ast.Name) and isinstance(ast_node.node.value, ast.Call):
118 node_id = n.id
119 func = ast_node.node.value.func
120 ref_id = func.id if isinstance(func, ast.Name) else None
121 if ref_id:
122 ref_kind = "CallRef"
123 self.add_reference(node_id, ref_id, ref_kind)
124 case ast.AnnAssign:
125 if isinstance(ast_node.node, ast.AnnAssign) and (
126 ast_node.node.annotation
127 and isinstance(ast_node.node.target, ast.Name)
128 and isinstance(ast_node.node.annotation, ast.Name)
129 ):
130 node_id = ast_node.node.target.id
131 ref_id = ast_node.node.annotation.id
132 ref_kind = "TypeRef"
133 self.add_reference(node_id, ref_id, ref_kind)
134 case ast.ClassDef:
135 if isinstance(ast_node.node, ast.ClassDef):
136 node = ast_node.node
137 node_id = node.name
138 if node.bases:
139 ref_node = node.bases[0]
140 if isinstance(ref_node, ast.Name):
141 ref_id = ref_node.id
142 ref_kind = "Inherit"
143 self.add_reference(node_id, ref_id, ref_kind)
144 # add functions and attributes to class
146 case ast.Call:
147 if isinstance(ast_node.node, ast.Call):
148 # obj.function. then obj refers to function
149 if isinstance(ast_node.node.func, ast.Attribute):
150 node_id = ast_node.name
151 ref_id = ast_node.node.func.attr
152 ref_kind = "FuncCall"
153 self.add_reference(node_id, ref_id, ref_kind)
154 # call function 'a' in function 'b', then 'b' refers to 'a'
155 container = ast_node.get_container_parent()
156 if container.ast_type == FunctionDef and isinstance(ast_node.node.func, ast.Name):
157 node_id = container.name
158 ref_id = ast_node.node.func.id
159 ref_kind = "FuncCall"
160 self.add_reference(node_id, ref_id, ref_kind)
162 def add_reference(self, node_id: str, ref_id: str, ref_kind: str) -> None:
163 properties = {}
164 if node_id == ref_id:
165 return
166 reference = PythonRSTReference(ref_id, ref_kind, properties)
167 referenced_by = PythonRSTReference(node_id, ref_kind, properties)
168 if node_id in self._references:
169 self._references[node_id].append(reference)
170 else:
171 self._references[node_id] = [reference]
172 if ref_id in self._referenced_by:
173 self._referenced_by[ref_id].append(referenced_by)
174 else:
175 self._referenced_by[ref_id] = [referenced_by]
177 def get_referenced_by(self, node_id):
178 refs = self._referenced_by.get(node_id, [])
179 return [PythonRSTReference(self._nodes[ref.node_id].name, ref.ref_kind, ref.properties) for ref in refs]
181 def get_references(self, node_id):
182 refs = self._references.get(node_id, [])
183 return [PythonRSTReference(self._nodes[ref.node_id].name, ref.ref_kind, ref.properties) for ref in refs]
186class PythonRstNode:
187 def __init__(self, node: ast.AST, translation_unit: PythonRstTranslationUnit = None, parent=None):
188 self.root = parent.root if parent and parent.root else self
189 self.node = node
190 self.parent = parent
191 self.translation_unit: PythonRstTranslationUnit = translation_unit
192 self.ast_type = KIND_MAP.get(type(node).__name__, UnknownType)
193 if self.ast_type == UnknownType:
194 print(f'"{type(node).__name__}": {type(node).__name__},')
196 self.indent = ""
197 self.name = self._derive_name()
198 self.show_props = False
199 self.children = []
200 self.properties = {}
201 self.is_implicit = self.ast_type not in IMPLICIT
202 self.offset = 0
203 self.length = 0
204 if self.translation_unit:
205 self.filename = translation_unit.file_name
206 self.derive_position(node, translation_unit, parent)
207 self.add_node()
208 for name in node._fields:
209 try:
210 child = getattr(node, name)
211 match child:
212 case list(): # Matches any list
213 if isinstance(node, Global) and name == "names":
214 if len(child) == 1:
215 self.name = child[0]
216 if name == "body":
217 self.body = self.children
219 if isinstance(node, (ImplicitNode, ast.Module)) or len(node._fields) == 1:
220 self.children.extend(PythonRstNode(n, translation_unit, self) for n in child)
221 if name == "body":
222 self.body = self.children
223 else:
224 self.children.append(PythonRstNode(ImplicitNode(name, child), translation_unit, self))
225 if name in ["body", "cases"]:
226 self.body = self.children[-1].children
228 case ast.AST():
229 if name not in ["ctx"]:
230 self.children.append(PythonRstNode(child, translation_unit, self))
231 if isinstance(child, ast.expr):
232 self.expression = self.children[-1]
233 case _:
234 if name not in ["None"]:
235 self.properties[name] = child
236 except AttributeError as e:
237 print(e)
238 continue
240 self.end_offset = self.offset + self.length
241 self.extended_end_offset = self.end_offset
242 self.is_statement = isinstance(self.node, ast.stmt)
244 def __eq__(self, other):
245 return (
246 isinstance(other, type(self))
247 and self.ast_type == other.ast_type
248 and match_props(self.properties, other.properties, IRRELEVANT_PROPS)
249 and match_children(self.children, other.children, IRRELEVANT_NODES)
250 )
252 def __contains__(self, item):
253 if not isinstance(item, list):
254 item = [item]
255 return find_in_list(self.children, item)
257 def __getitem__(self, key):
258 """Allow indexing/slicing into node to access children.
260 Usage: node[0] == node.children[0]
261 """
262 return self.children[key]
264 def __repr__(self):
265 return format_node(self)
267 @property
268 def next_sibling(self) -> Self | None:
269 return next_sibling(self)
271 @property
272 def preceding_sibling(self) -> Self | None:
273 return preceding_sibling(self)
275 def process(self, function: Callable[[Self], None]) -> None:
276 function(self)
277 for child in self.children:
278 child.process(function)
280 def derive_position(self, node: ast.AST, translation_unit: PythonRstTranslationUnit, parent):
281 if node._attributes:
282 if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)) and node.decorator_list:
283 self.offset = convert(self.translation_unit.lines, node.decorator_list[0].lineno, node.decorator_list[0].col_offset) - 1
284 elif parent.name == "decorator_list":
285 # also include the @ in the decorator
286 self.offset = convert(self.translation_unit.lines, node.lineno, node.col_offset) - 1 # type: ignore[attr-defined]
287 else:
288 self.offset = convert(self.translation_unit.lines, node.lineno, node.col_offset) # type: ignore[attr-defined]
289 all_space = all(c == " " for c in self.translation_unit.content[self.offset - node.col_offset : self.offset])
290 if all_space:
291 self.offset = max(self.offset - node.col_offset, 0)
292 self.length = convert(self.translation_unit.lines, node.end_lineno, node.end_col_offset) - self.offset # type: ignore[attr-defined]
293 elif isinstance(node, ast.Module) and translation_unit:
294 self.offset = 0
295 self.length = len(translation_unit.content)
296 else:
297 self.offset = 0
298 self.length = 0
300 @staticmethod
301 def load(
302 file_path: Path, # TODO: why Path - why not FileDescriptorOrPath (the type of the file parameter of the open function)?
303 extra_args: Sequence[str] | None = None,
304 working_dir: Path | None = None,
305 ) -> PythonRstNode:
306 # Keep a uniform loader signature across AST node implementations.
307 # Python's AST parser does not need extra arguments or a working dir.
308 _ = extra_args, working_dir
309 with Path(file_path).open() as file:
310 content = file.read()
311 return PythonRstNode.load_from_text(content, str(file_path))
313 @staticmethod
314 def load_from_text(
315 text: str,
316 file_name: str = "test.py",
317 extra_args: Sequence[str] | None = None,
318 working_dir: Path | None = None,
319 ) -> PythonRstNode:
320 _ = extra_args, working_dir
321 translation_unit = PythonRstTranslationUnit(text, file_name=str(file_name))
322 translation_unit.check_diagnostics()
323 root_node = PythonRstNode(translation_unit.atu, translation_unit)
324 return root_node
326 def _derive_name(self):
328 if (
329 isinstance(
330 self.node,
331 (
332 ast.FunctionDef,
333 ast.AsyncFunctionDef,
334 ast.ClassDef,
335 ast.ExceptHandler,
336 ),
337 )
338 and self.node.name
339 ):
340 name = self.node.name
341 elif isinstance(self.node, ast.Global) and len(self.node.names) == 1:
342 name = self.node.names[0]
343 elif isinstance(self.node, (ast.AnnAssign, ast.AugAssign)) and isinstance(self.node.target, ast.Name):
344 name = self.node.target.id
345 elif isinstance(self.node, ast.Assign) and len(self.node.targets) == 1:
346 target = self.node.targets[0]
347 name = target.id if isinstance(target, ast.Name) else self.ast_type.__name__
348 elif isinstance(self.node, ast.Name):
349 name = self.node.id
350 elif isinstance(self.node, ast.arg):
351 name = self.node.arg
352 elif isinstance(self.node, ast.Match) and isinstance(self.node.subject, ast.Name):
353 name = self.node.subject.id
354 elif (isinstance(self.node, ast.Import) and len(self.node.names) == 1) or (
355 isinstance(self.node, ast.ImportFrom) and len(self.node.names) == 1
356 ):
357 name = self.node.names[0].name
358 elif isinstance(self.node, (ast.Assert, ast.Break, ast.Pass, ast.Raise, ast.Continue)):
359 name = ""
360 elif isinstance(self.node, (ast.For, ast.AsyncFor)):
361 if isinstance(self.node.target, Tuple):
362 name = self.node.target.dims[1].id
363 elif isinstance(self.node.target, ast.Name):
364 name = self.node.target.id
365 else:
366 name = str(self.node.target)
367 elif "body" not in self.node._fields:
368 name = ast.unparse(self.node)
369 elif isinstance(self.node, (ast.Module)) and self.translation_unit:
370 name = self.translation_unit.file_name
371 else:
372 name = self.ast_type.__name__
373 return name or ""
375 @property
376 def type(self):
377 return self.node.annotation.id if isinstance(self.node, ast.AnnAssign) and isinstance(self.node.annotation, ast.Name) else None
379 @property
380 def value(self):
381 if self.ast_type == Assert:
382 return 0
383 return self.node.value.value if hasattr(self.node, "value") else None
385 @property
386 def expr(self):
387 if (
388 isinstance(
389 self.node,
390 (
391 ast.Assign,
392 ast.AnnAssign,
393 ast.AugAssign,
394 ast.Return,
395 ast.Expr,
396 ast.Delete,
397 ast.NamedExpr,
398 ),
399 )
400 and hasattr(self.node, "value")
401 and self.node.value is not None
402 ) or (isinstance(self.node, ast.Expr) and hasattr(self.node, "value")):
403 return PythonRstNode(self.node.value, self.translation_unit, self)
404 if isinstance(self.node, (ast.For, ast.AsyncFor, ast.comprehension)):
405 return PythonRstNode(self.node.iter, self.translation_unit, self)
406 if isinstance(self.node, (ast.If, ast.While, ast.Assert)):
407 return PythonRstNode(self.node.test, self.translation_unit, self)
408 if isinstance(self.node, (ast.Raise, ast.ExceptHandler)) and hasattr(self.node, "exc") and self.node.exc is not None:
409 return PythonRstNode(self.node.exc, self.translation_unit, self)
410 return None
412 @property
413 def operator(self):
414 node_type = type(self.node).__name__
415 op = type(self.node.op).__name__ if isinstance(self.node, (ast.BinOp, ast.UnaryOp, ast.BoolOp, ast.AugAssign)) else ""
416 return OPERATOR_MAP.get(node_type + op, "")
418 @property
419 def signature(self) -> str:
420 sig = self.binary_file_content().decode(sys.getfilesystemencoding())
421 if self.parent and self.parent.name == "decorator_list" and not sig.startswith("@"):
422 sig = "@" + sig
423 return sig
425 def binary_file_content(self) -> bytes:
426 return (
427 self.translation_unit.content[self.offset : self.offset + self.length]
428 if self.translation_unit
429 else ast.unparse(self.node).encode(sys.getfilesystemencoding())
430 )
432 @property
433 def referenced_by(self) -> Sequence[PythonRSTReference]:
434 self.translation_unit.lazy_create_refers(self)
435 return self.translation_unit.get_referenced_by(self.name)
437 @property
438 def references(self) -> list[PythonRSTReference]:
439 self.translation_unit.lazy_create_refers(self)
440 return self.translation_unit.get_references(self.name)
442 def add_node(self):
443 self.translation_unit.add(self)
445 def get_container_parent(self):
446 if self.parent:
447 if self.parent.ast_type.__name__ in ["FunctionDef", "ClassDef", "Module"]:
448 return self.parent
449 return self.parent.get_container_parent()
450 return self
452 @property
453 def text(self) -> str:
454 return textwrap.dedent(self.signature)