Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/astroid/transforms.py: 43%
Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1# Licensed under the LGPL: https://www.gnu.org/licenses/old-licenses/lgpl-2.1.en.html
2# For details: https://github.com/pylint-dev/astroid/blob/main/LICENSE
3# Copyright (c) https://github.com/pylint-dev/astroid/blob/main/CONTRIBUTORS.txt
5from __future__ import annotations
7import warnings
8from collections import defaultdict
9from collections.abc import Callable
10from typing import TYPE_CHECKING, TypeVar, Union, cast, overload
12from astroid.context import _invalidate_cache
13from astroid.typing import SuccessfulInferenceResult, TransformFn
15if TYPE_CHECKING:
16 from astroid import nodes
18 _SuccessfulInferenceResultT = TypeVar(
19 "_SuccessfulInferenceResultT", bound=SuccessfulInferenceResult
20 )
21 _Predicate = Callable[[_SuccessfulInferenceResultT], bool] | None
23# pylint: disable-next=consider-alternative-union-syntax
24_Vistables = Union[
25 "nodes.NodeNG", list["nodes.NodeNG"], tuple["nodes.NodeNG", ...], str, None
26]
27_VisitReturns = (
28 SuccessfulInferenceResult
29 | list[SuccessfulInferenceResult]
30 | tuple[SuccessfulInferenceResult, ...]
31 | str
32 | None
33)
36class TransformVisitor:
37 """A visitor for handling transforms.
39 The standard approach of using it is to call
40 :meth:`~visit` with an *astroid* module and the class
41 will take care of the rest, walking the tree and running the
42 transforms for each encountered node.
44 Based on its usage in AstroidManager.brain, it should not be reinstantiated.
45 """
47 def __init__(self) -> None:
48 # The typing here is incorrect, but it's the best we can do
49 # Refer to register_transform and unregister_transform for the correct types
50 self.transforms: defaultdict[
51 type[SuccessfulInferenceResult],
52 list[
53 tuple[
54 TransformFn[SuccessfulInferenceResult],
55 _Predicate[SuccessfulInferenceResult],
56 ]
57 ],
58 ] = defaultdict(list)
60 def _transform(self, node: SuccessfulInferenceResult) -> SuccessfulInferenceResult:
61 """Call matching transforms for the given node if any and return the
62 transformed node.
63 """
64 cls = node.__class__
66 for transform_func, predicate in self.transforms[cls]:
67 if predicate is None or predicate(node):
68 ret = transform_func(node)
69 # if the transformation function returns something, it's
70 # expected to be a replacement for the node
71 if ret is not None:
72 _invalidate_cache()
73 node = ret
74 if ret.__class__ != cls:
75 # Can no longer apply the rest of the transforms.
76 break
77 return node
79 def _visit(self, node: nodes.NodeNG) -> SuccessfulInferenceResult:
80 for name in node._astroid_fields:
81 value = getattr(node, name)
82 if TYPE_CHECKING:
83 value = cast(_Vistables, value)
85 visited = self._visit_generic(value)
86 if visited != value:
87 setattr(node, name, visited)
88 return self._transform(node)
90 @overload
91 def _visit_generic(self, node: None) -> None: ...
93 @overload
94 def _visit_generic(self, node: str) -> str: ...
96 @overload
97 def _visit_generic(
98 self, node: list[nodes.NodeNG]
99 ) -> list[SuccessfulInferenceResult]: ...
101 @overload
102 def _visit_generic(
103 self, node: tuple[nodes.NodeNG, ...]
104 ) -> tuple[SuccessfulInferenceResult, ...]: ...
106 @overload
107 def _visit_generic(self, node: nodes.NodeNG) -> SuccessfulInferenceResult: ...
109 def _visit_generic(self, node: _Vistables) -> _VisitReturns:
110 if not node:
111 return node
112 if isinstance(node, list):
113 return [self._visit_generic(child) for child in node]
114 if isinstance(node, tuple):
115 return tuple(self._visit_generic(child) for child in node)
116 if isinstance(node, str):
117 return node
119 try:
120 return self._visit(node)
121 except RecursionError:
122 # Returning the node untransformed is better than giving up.
123 warnings.warn(
124 f"Astroid was unable to transform {node}.\n"
125 "Some functionality will be missing unless the system recursion limit is lifted.\n"
126 "From pylint, try: --init-hook='import sys; sys.setrecursionlimit(2000)' or higher.",
127 UserWarning,
128 stacklevel=0,
129 )
130 return node
132 def register_transform(
133 self,
134 node_class: type[_SuccessfulInferenceResultT],
135 transform: TransformFn[_SuccessfulInferenceResultT],
136 predicate: _Predicate[_SuccessfulInferenceResultT] | None = None,
137 ) -> None:
138 """Register `transform(node)` function to be applied on the given node.
140 The transform will only be applied if `predicate` is None or returns true
141 when called with the node as argument.
143 The transform function may return a value which is then used to
144 substitute the original node in the tree.
145 """
146 self.transforms[node_class].append((transform, predicate)) # type: ignore[index, arg-type]
148 def unregister_transform(
149 self,
150 node_class: type[_SuccessfulInferenceResultT],
151 transform: TransformFn[_SuccessfulInferenceResultT],
152 predicate: _Predicate[_SuccessfulInferenceResultT] | None = None,
153 ) -> None:
154 """Unregister the given transform."""
155 self.transforms[node_class].remove((transform, predicate)) # type: ignore[index, arg-type]
157 def visit(self, node: nodes.NodeNG) -> SuccessfulInferenceResult:
158 """Walk the given astroid *tree* and transform each encountered node.
160 Only the nodes which have transforms registered will actually
161 be replaced or changed.
162 """
163 return self._visit(node)