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

67 statements  

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 

4 

5from __future__ import annotations 

6 

7import warnings 

8from collections import defaultdict 

9from collections.abc import Callable 

10from typing import TYPE_CHECKING, TypeVar, Union, cast, overload 

11 

12from astroid.context import _invalidate_cache 

13from astroid.typing import SuccessfulInferenceResult, TransformFn 

14 

15if TYPE_CHECKING: 

16 from astroid import nodes 

17 

18 _SuccessfulInferenceResultT = TypeVar( 

19 "_SuccessfulInferenceResultT", bound=SuccessfulInferenceResult 

20 ) 

21 _Predicate = Callable[[_SuccessfulInferenceResultT], bool] | None 

22 

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) 

34 

35 

36class TransformVisitor: 

37 """A visitor for handling transforms. 

38 

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. 

43 

44 Based on its usage in AstroidManager.brain, it should not be reinstantiated. 

45 """ 

46 

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) 

59 

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__ 

65 

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 

78 

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) 

84 

85 visited = self._visit_generic(value) 

86 if visited != value: 

87 setattr(node, name, visited) 

88 return self._transform(node) 

89 

90 @overload 

91 def _visit_generic(self, node: None) -> None: ... 

92 

93 @overload 

94 def _visit_generic(self, node: str) -> str: ... 

95 

96 @overload 

97 def _visit_generic( 

98 self, node: list[nodes.NodeNG] 

99 ) -> list[SuccessfulInferenceResult]: ... 

100 

101 @overload 

102 def _visit_generic( 

103 self, node: tuple[nodes.NodeNG, ...] 

104 ) -> tuple[SuccessfulInferenceResult, ...]: ... 

105 

106 @overload 

107 def _visit_generic(self, node: nodes.NodeNG) -> SuccessfulInferenceResult: ... 

108 

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 

118 

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 

131 

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. 

139 

140 The transform will only be applied if `predicate` is None or returns true 

141 when called with the node as argument. 

142 

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] 

147 

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] 

156 

157 def visit(self, node: nodes.NodeNG) -> SuccessfulInferenceResult: 

158 """Walk the given astroid *tree* and transform each encountered node. 

159 

160 Only the nodes which have transforms registered will actually 

161 be replaced or changed. 

162 """ 

163 return self._visit(node)