Coverage for src/renaissance/integrations/python/ast/rst_node.py: 96%
332 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-15 10:34 +0000
1# Copyright (C) 2024-2026 TNO
2#
3# Licensed under the terms specified in the LICENSE file in the root directory
4# of this repository, also available at:
5# https://github.com/TNO/Renaissance.Py/blob/main/LICENSE
7"""AI: ASTNode implementation backed by Python's built-in `ast` module (the "RST" representation)."""
9import ast
10import sys
11import textwrap
12from collections.abc import Callable, Sequence
13from pathlib import Path
14from typing import Any, Self
16from renaissance.integrations.python.ast.kinds import PYTHON_KIND_MAP, PYTHON_OPERATOR_MAP
17from renaissance.integrations.python.ast.util import convert
18from renaissance.syntax_tree.match_finder import find_in_list
19from renaissance.syntax_tree.semantic_kind import SemanticKind
20from renaissance.utils.ast_utils import (
21 format_node,
22 match_children,
23 match_props,
24 next_sibling,
25 preceding_sibling,
26 traverse,
27)
29types = ["int", "float", "str", "list", "set", "tuple", "Mapping", "dict", "Optional"]
30IRRELEVANT_PROPS = {"comment"}
31IRRELEVANT_NODES = {"comment"}
32IMPLICIT = {"ImplicitNode"}
35class ImplicitNode(ast.Name):
36 """AI: Represent a synthetic AST node inserted where the real source has none."""
38 _fields = (
39 "id",
40 "body",
41 )
43 _field_types = {
44 "id": str,
45 "body": list,
46 }
48 def __init__(self, name, children=None) -> None:
49 """AI: Represent a synthetic AST node inserted where the real source has none."""
50 super().__init__(name, children or [])
51 self.lineno = 0
52 self.col_offset = 0
53 self.end_lineno = 0
54 self.end_col_offset = 0
57class PythonRSTReference:
58 """AI: Represent a reference from one Python RST AST node to another by id."""
60 def __repr__(self) -> str:
61 """AI: Return a string identifying the referenced node id and reference kind."""
62 return f"{self.node_id}:{self.ref_kind}"
64 def __init__(self, node_id: str, ref_kind: str, properties: dict[str, Any]) -> None:
65 """AI: Represent a reference from one Python RST AST node to another by id."""
66 self.node_id = node_id
67 self.ref_kind = ref_kind
68 self.properties = properties
71class PythonRstTranslationUnit:
72 """AI: Parse Python source into a stdlib ast tree with lazily-built reference caches."""
74 cache = {}
76 def __init__(self, content, file_name: str) -> None:
77 """AI: Parse Python source into a stdlib ast tree with lazily-built reference caches."""
78 # TODO: see docs/developer/architecture/position-consistency.md - PEP 263 declares the encoding, not the host.
79 self.content = content.encode(sys.getfilesystemencoding())
80 self.atu = ast.parse(content, file_name)
81 self.file_name = file_name
82 self.references_initialized = False
83 PythonRstTranslationUnit.cache[file_name] = content
84 self.lines = self.content.splitlines()
86 self._references: dict[str, list[PythonRSTReference]] = {}
87 self._referenced_by: dict[str, list[PythonRSTReference]] = {}
88 self._nodes: dict[str, PythonRstNode] = {}
90 def check_diagnostics(self, continue_with_warning=True) -> None:
91 """AI: Report ast type-ignore comments, raising an Exception unless continue_with_warning is True."""
92 msg = None
93 errors = ""
94 for d in self.atu.type_ignores:
95 msg = f"type ignored: {d.tag} at {d.lineno}\n"
96 errors += msg
97 print(msg)
98 if msg and not continue_with_warning:
99 raise Exception(f"Error parsing: {self.file_name} \n+ errors: {errors}")
101 def lazy_create_refers(self, node: PythonRstNode) -> None:
102 """AI: Build the translation unit's reference cache on first use, then no-op on subsequent calls."""
103 if self.references_initialized:
104 return
105 for n in traverse(node.root):
106 self.create_references(n)
107 self.references_initialized = True
109 def add(self, node) -> None:
110 """AI: Register node in the node-name lookup table, keyed by its parser-kind-specific name."""
111 match node.parser_kind:
112 case "Name":
113 if node.node.id not in self._nodes and node.node.id not in types:
114 self._nodes[node.node.id] = node
115 case "FunctionDef":
116 if node.node.name not in self._nodes:
117 self._nodes[node.node.name] = node
118 case "Call":
119 if node.name not in self._nodes:
120 self._nodes[node.name] = node
121 case "ClassDef":
122 if node.name not in self._nodes:
123 self._nodes[node.name] = node
124 case "arg":
125 if node.name != "self" and node.name not in self._nodes:
126 self._nodes[node.name] = node
128 def create_references(self, ast_node) -> None:
129 """AI: Derive and record type, call, and inheritance references for ast_node based on its ast type."""
130 assert isinstance(ast_node, PythonRstNode), f"Expected PythonASTNode but got {type(ast_node)}"
131 match type(ast_node.node):
132 case ast.arg:
133 # TODO: is excluding self here intentional..?
134 if ast_node.name != "self" and isinstance(ast_node.node, ast.arg) and isinstance(ast_node.node.annotation, ast.Name):
135 node_id = ast_node.name
136 ref_id = ast_node.node.annotation.id
137 ref_kind = "TypeRef"
138 self.add_reference(node_id, ref_id, ref_kind)
139 case ast.FunctionDef | ast.AsyncFunctionDef:
140 if isinstance(ast_node.node, (ast.FunctionDef, ast.AsyncFunctionDef)) and isinstance(ast_node.node.returns, ast.Name):
141 node_id = ast_node.name
142 ref_id = ast_node.node.returns.id
143 ref_kind = "TypeRef"
144 self.add_reference(node_id, ref_id, ref_kind)
145 case ast.Assign:
146 if isinstance(ast_node.node, ast.Assign):
147 for n in ast_node.node.targets:
148 if isinstance(n, ast.Name) and isinstance(ast_node.node.value, ast.Call):
149 node_id = n.id
150 func = ast_node.node.value.func
151 ref_id = func.id if isinstance(func, ast.Name) else None
152 if ref_id:
153 ref_kind = "CallRef"
154 self.add_reference(node_id, ref_id, ref_kind)
155 case ast.AnnAssign:
156 if isinstance(ast_node.node, ast.AnnAssign) and (
157 ast_node.node.annotation
158 and isinstance(ast_node.node.target, ast.Name)
159 and isinstance(ast_node.node.annotation, ast.Name)
160 ):
161 node_id = ast_node.node.target.id
162 ref_id = ast_node.node.annotation.id
163 ref_kind = "TypeRef"
164 self.add_reference(node_id, ref_id, ref_kind)
165 case ast.ClassDef:
166 if isinstance(ast_node.node, ast.ClassDef):
167 node = ast_node.node
168 node_id = node.name
169 if node.bases:
170 ref_node = node.bases[0]
171 if isinstance(ref_node, ast.Name):
172 ref_id = ref_node.id
173 ref_kind = "Inherit"
174 self.add_reference(node_id, ref_id, ref_kind)
175 # add functions and attributes to class
177 case ast.Call:
178 if isinstance(ast_node.node, ast.Call):
179 # obj.function. then obj refers to function
180 if isinstance(ast_node.node.func, ast.Attribute):
181 node_id = ast_node.name
182 ref_id = ast_node.node.func.attr
183 ref_kind = "FuncCall"
184 self.add_reference(node_id, ref_id, ref_kind)
185 # call function 'a' in function 'b', then 'b' refers to 'a'
186 container = ast_node.get_container_parent()
187 if container.parser_kind in {"FunctionDef", "AsyncFunctionDef"} and isinstance(ast_node.node.func, ast.Name):
188 node_id = container.name
189 ref_id = ast_node.node.func.id
190 ref_kind = "FuncCall"
191 self.add_reference(node_id, ref_id, ref_kind)
193 def add_reference(self, node_id: str, ref_id: str, ref_kind: str) -> None:
194 """AI: Record a reference from node_id to ref_id (and the corresponding referenced-by entry)."""
195 properties = {}
196 if node_id == ref_id:
197 return
198 reference = PythonRSTReference(ref_id, ref_kind, properties)
199 referenced_by = PythonRSTReference(node_id, ref_kind, properties)
200 if node_id in self._references:
201 self._references[node_id].append(reference)
202 else:
203 self._references[node_id] = [reference]
204 if ref_id in self._referenced_by:
205 self._referenced_by[ref_id].append(referenced_by)
206 else:
207 self._referenced_by[ref_id] = [referenced_by]
209 def get_referenced_by(self, node_id):
210 """AI: Return the references pointing to node_id, resolved to their referring nodes' names."""
211 refs = self._referenced_by.get(node_id, [])
212 return [PythonRSTReference(self._nodes[ref.node_id].name, ref.ref_kind, ref.properties) for ref in refs]
214 def get_references(self, node_id):
215 """AI: Return the references that node_id points to, resolved to their target nodes' names."""
216 refs = self._references.get(node_id, [])
217 return [PythonRSTReference(self._nodes[ref.node_id].name, ref.ref_kind, ref.properties) for ref in refs]
220class PythonRstNode:
221 """AI: ASTNode implementation backed by the stdlib ast module."""
223 def __init__(self, node: ast.AST, translation_unit: PythonRstTranslationUnit = None, parent=None) -> None:
224 """AI: Wrap a stdlib ast node as an AST node within the given translation unit."""
225 self.root = parent.root if parent and parent.root else self
226 self.node = node
227 self.parent = parent
228 self.translation_unit: PythonRstTranslationUnit = translation_unit
229 self.parser_kind = type(node).__name__
230 self.semantic_kind = PYTHON_KIND_MAP.get(self.parser_kind, SemanticKind.NODE)
232 self.indent = ""
233 self.name = self._derive_name()
234 self.show_props = False
235 self.children = []
236 self.properties = {}
237 self.is_implicit = self.parser_kind not in IMPLICIT
238 self.offset = 0
239 self.length = 0
240 if self.translation_unit:
241 self.filename = translation_unit.file_name
242 self.derive_position(node, translation_unit, parent)
243 self.add_node()
244 for name in node._fields:
245 try:
246 child = getattr(node, name)
247 match child:
248 case list(): # Matches any list
249 if isinstance(node, ast.Global) and name == "names" and len(child) == 1:
250 self.name = child[0]
252 if isinstance(node, (ast.Global, ast.Nonlocal)):
253 pass # names is list[str], not AST nodes - nothing to build children from
254 elif isinstance(node, (ImplicitNode, ast.Module)) or len(node._fields) == 1:
255 # A list field can hold a bare None at a position with no value
256 # None isn't a real AST node, so it has nothing to build a child from.
257 self.children.extend(PythonRstNode(n, translation_unit, self) for n in child if n is not None)
258 if name == "body":
259 self.body = self.children
260 else:
261 self.children.append(PythonRstNode(ImplicitNode(name, child), translation_unit, self))
262 if name in ["body", "cases"]:
263 self.body = self.children[-1].children
265 case ast.AST():
266 if name != "ctx":
267 self.children.append(PythonRstNode(child, translation_unit, self))
268 if isinstance(child, ast.expr):
269 self.expression = self.children[-1]
270 case _:
271 if name != "None":
272 self.properties[name] = child
273 except AttributeError as e:
274 print(e)
275 continue
277 self.end_offset = self.offset + self.length
278 self.extended_end_offset = self.end_offset
279 self.is_statement = isinstance(self.node, ast.stmt)
281 @property
282 def kind_key(self) -> SemanticKind | str:
283 """AI: Return the semantic kind, or the raw parser kind when no semantic kind applies."""
284 return self.semantic_kind if self.semantic_kind is not SemanticKind.NODE else self.parser_kind
286 def __eq__(self, other) -> bool:
287 """AI: Return whether this node is structurally equal to `other`, ignoring irrelevant properties/children."""
288 return (
289 isinstance(other, type(self))
290 and self.kind_key == other.kind_key
291 and match_props(self.properties, other.properties, IRRELEVANT_PROPS)
292 and match_children(self.children, other.children, IRRELEVANT_NODES)
293 )
295 def __contains__(self, item) -> bool:
296 """AI: Return whether item(s) are found among this node's children."""
297 if not isinstance(item, list):
298 item = [item]
299 return find_in_list(self.children, item) >= 0
301 def __getitem__(self, key) -> Any: # noqa: ANN401 (an int key yields a node, a slice yields a list of nodes)
302 """Allow indexing/slicing into node to access children.
304 Usage: node[0] == node.children[0]
305 """
306 return self.children[key]
308 def __repr__(self) -> str:
309 """AI: Return the formatted node representation."""
310 return format_node(self)
312 @property
313 def next_sibling(self) -> Self | None:
314 """AI: Return the sibling node immediately following this one, or None."""
315 return next_sibling(self)
317 @property
318 def preceding_sibling(self) -> Self | None:
319 """AI: Return the sibling node immediately preceding this one, or None."""
320 return preceding_sibling(self)
322 def process(self, function: Callable[[Self], None]) -> None:
323 """AI: Apply function to this node and recursively to all of its descendants."""
324 function(self)
325 for child in self.children:
326 child.process(function)
328 def derive_position(self, node: ast.AST, translation_unit: PythonRstTranslationUnit, parent) -> None:
329 """AI: Compute and set this node's offset and length in the source text from its ast position attributes."""
330 if node._attributes:
331 if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)) and node.decorator_list:
332 self.offset = convert(self.translation_unit.lines, node.decorator_list[0].lineno, node.decorator_list[0].col_offset) - 1
333 elif parent.name == "decorator_list":
334 # also include the @ in the decorator
335 self.offset = convert(self.translation_unit.lines, node.lineno, node.col_offset) - 1 # type: ignore[attr-defined]
336 else:
337 self.offset = convert(self.translation_unit.lines, node.lineno, node.col_offset) # type: ignore[attr-defined]
338 all_space = all(c == " " for c in self.translation_unit.content[self.offset - node.col_offset : self.offset])
339 if all_space:
340 self.offset = max(self.offset - node.col_offset, 0)
341 self.length = convert(self.translation_unit.lines, node.end_lineno, node.end_col_offset) - self.offset # type: ignore[attr-defined]
342 elif isinstance(node, ast.Module) and translation_unit:
343 self.offset = 0
344 self.length = len(translation_unit.content)
345 else:
346 self.offset = 0
347 self.length = 0
349 @staticmethod
350 def load(
351 file_path: Path, # TODO: why Path - why not FileDescriptorOrPath (the type of the file parameter of the open function)?
352 extra_args: Sequence[str] | None = None,
353 working_dir: Path | None = None,
354 ) -> PythonRstNode:
355 """AI: Parse the Python source file at file_path into a PythonRstNode tree."""
356 # Keep a uniform loader signature across AST node implementations.
357 # Python's AST parser does not need extra arguments or a working dir.
358 _ = extra_args, working_dir
359 with Path(file_path).open() as file:
360 content = file.read()
361 return PythonRstNode.load_from_text(content, str(file_path))
363 @staticmethod
364 def load_from_text(
365 text: str,
366 file_name: str = "test.py",
367 extra_args: Sequence[str] | None = None,
368 working_dir: Path | None = None,
369 ) -> PythonRstNode:
370 """AI: Parse Python source text into a PythonRstNode tree, attributed to file_name."""
371 _ = extra_args, working_dir
372 translation_unit = PythonRstTranslationUnit(text, file_name=str(file_name))
373 translation_unit.check_diagnostics()
374 root_node = PythonRstNode(translation_unit.atu, translation_unit)
375 return root_node
377 def _derive_name(self):
379 if (
380 isinstance(
381 self.node,
382 (
383 ast.FunctionDef,
384 ast.AsyncFunctionDef,
385 ast.ClassDef,
386 ast.ExceptHandler,
387 ),
388 )
389 and self.node.name
390 ):
391 name = self.node.name
392 elif isinstance(self.node, ast.Global) and len(self.node.names) == 1:
393 name = self.node.names[0]
394 elif isinstance(self.node, (ast.AnnAssign, ast.AugAssign)) and isinstance(self.node.target, ast.Name):
395 name = self.node.target.id
396 elif isinstance(self.node, ast.Assign) and len(self.node.targets) == 1:
397 target = self.node.targets[0]
398 name = target.id if isinstance(target, ast.Name) else self.parser_kind
399 elif isinstance(self.node, ast.Name):
400 name = self.node.id
401 elif isinstance(self.node, ast.arg):
402 name = self.node.arg
403 elif isinstance(self.node, ast.Match) and isinstance(self.node.subject, ast.Name):
404 name = self.node.subject.id
405 elif (isinstance(self.node, ast.Import) and len(self.node.names) == 1) or (
406 isinstance(self.node, ast.ImportFrom) and len(self.node.names) == 1
407 ):
408 name = self.node.names[0].name
409 elif isinstance(self.node, (ast.Assert, ast.Break, ast.Pass, ast.Raise, ast.Continue)):
410 name = ""
411 elif isinstance(self.node, (ast.For, ast.AsyncFor)):
412 if (
413 isinstance(self.node.target, ast.Tuple)
414 and len(self.node.target.elts) > 1
415 and isinstance(self.node.target.elts[1], ast.Name)
416 ):
417 name = self.node.target.elts[1].id
418 elif isinstance(self.node.target, ast.Name):
419 name = self.node.target.id
420 else:
421 name = str(self.node.target)
422 elif "body" not in self.node._fields:
423 name = ast.unparse(self.node)
424 elif isinstance(self.node, (ast.Module)) and self.translation_unit:
425 name = self.translation_unit.file_name
426 else:
427 name = self.parser_kind
428 return name or ""
430 @property
431 def type(self):
432 """AI: Return the annotation name for an annotated assignment node, or None otherwise."""
433 return self.node.annotation.id if isinstance(self.node, ast.AnnAssign) and isinstance(self.node.annotation, ast.Name) else None
435 @property
436 def value(self):
437 """AI: Return the literal value of this node, or None if it has none."""
438 if self.parser_kind == "Assert":
439 return 0
440 return self.node.value.value if hasattr(self.node, "value") else None
442 @property
443 def expr(self):
444 """AI: Return the wrapped inner expression node relevant to this statement's kind, or None."""
445 if (
446 isinstance(
447 self.node,
448 (
449 ast.Assign,
450 ast.AnnAssign,
451 ast.AugAssign,
452 ast.Return,
453 ast.Expr,
454 ast.Delete,
455 ast.NamedExpr,
456 ),
457 )
458 and hasattr(self.node, "value")
459 and self.node.value is not None
460 ) or (isinstance(self.node, ast.Expr) and hasattr(self.node, "value")):
461 return PythonRstNode(self.node.value, self.translation_unit, self)
462 if isinstance(self.node, (ast.For, ast.AsyncFor, ast.comprehension)):
463 return PythonRstNode(self.node.iter, self.translation_unit, self)
464 if isinstance(self.node, (ast.If, ast.While, ast.Assert)):
465 return PythonRstNode(self.node.test, self.translation_unit, self)
466 if isinstance(self.node, (ast.Raise, ast.ExceptHandler)) and hasattr(self.node, "exc") and self.node.exc is not None:
467 return PythonRstNode(self.node.exc, self.translation_unit, self)
468 return None
470 @property
471 def operator(self):
472 """AI: Return this node's operator symbol, or an empty string if it has none."""
473 node_type = type(self.node).__name__
474 op = type(self.node.op).__name__ if isinstance(self.node, (ast.BinOp, ast.UnaryOp, ast.BoolOp, ast.AugAssign)) else ""
475 return PYTHON_OPERATOR_MAP.get(node_type + op, "")
477 @property
478 def signature(self) -> str:
479 """AI: Return the source code text of this node, prefixed with '@' for decorators."""
480 # TODO: see docs/developer/architecture/position-consistency.md - decodes with the host encoding, not the adapter's.
481 sig = self.binary_file_content().decode(sys.getfilesystemencoding())
482 if self.parent and self.parent.name == "decorator_list" and not sig.startswith("@"):
483 sig = "@" + sig
484 return sig
486 def binary_file_content(self) -> bytes:
487 """AI: Return this node's source text as encoded bytes."""
488 # TODO: see docs/developer/architecture/position-consistency.md - the unparse fallback builds a buffer no offset indexes.
489 return (
490 self.translation_unit.content[self.offset : self.offset + self.length]
491 if self.translation_unit
492 else ast.unparse(self.node).encode(sys.getfilesystemencoding())
493 )
495 @property
496 def referenced_by(self) -> Sequence[PythonRSTReference]:
497 """AI: Return the references that point to this node, building the reference cache if needed."""
498 self.translation_unit.lazy_create_refers(self)
499 return self.translation_unit.get_referenced_by(self.name)
501 @property
502 def references(self) -> list[PythonRSTReference]:
503 """AI: Return the references that this node points to, building the reference cache if needed."""
504 self.translation_unit.lazy_create_refers(self)
505 return self.translation_unit.get_references(self.name)
507 def add_node(self) -> None:
508 """AI: Register this node in its translation unit's node-name lookup table."""
509 self.translation_unit.add(self)
511 def get_container_parent(self):
512 """AI: Return the nearest ancestor node that is a function, class, or module."""
513 if self.parent:
514 if self.parent.parser_kind in {"FunctionDef", "AsyncFunctionDef", "ClassDef", "Module"}:
515 return self.parent
516 return self.parent.get_container_parent()
517 return self
519 @property
520 def text(self) -> str:
521 """AI: Return this node's signature text with common leading whitespace removed."""
522 return textwrap.dedent(self.signature)