Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.10/site-packages/astroid/transforms.py: 76%

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, Optional, 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 = Optional[Callable[[_SuccessfulInferenceResultT], bool]] 

22 

23_Vistables = Union[ 

24 "nodes.NodeNG", list["nodes.NodeNG"], tuple["nodes.NodeNG", ...], str, None 

25] 

26_VisitReturns = Union[ 

27 SuccessfulInferenceResult, 

28 list[SuccessfulInferenceResult], 

29 tuple[SuccessfulInferenceResult, ...], 

30 str, 

31 None, 

32] 

33 

34 

35class TransformVisitor: 

36 """A visitor for handling transforms. 

37 

38 The standard approach of using it is to call 

39 :meth:`~visit` with an *astroid* module and the class 

40 will take care of the rest, walking the tree and running the 

41 transforms for each encountered node. 

42 

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

44 """ 

45 

46 def __init__(self) -> None: 

47 # The typing here is incorrect, but it's the best we can do 

48 # Refer to register_transform and unregister_transform for the correct types 

49 self.transforms: defaultdict[ 

50 type[SuccessfulInferenceResult], 

51 list[ 

52 tuple[ 

53 TransformFn[SuccessfulInferenceResult], 

54 _Predicate[SuccessfulInferenceResult], 

55 ] 

56 ], 

57 ] = defaultdict(list) 

58 

59 def _transform(self, node: SuccessfulInferenceResult) -> SuccessfulInferenceResult: 

60 """Call matching transforms for the given node if any and return the 

61 transformed node. 

62 """ 

63 cls = node.__class__ 

64 

65 for transform_func, predicate in self.transforms[cls]: 

66 if predicate is None or predicate(node): 

67 ret = transform_func(node) 

68 # if the transformation function returns something, it's 

69 # expected to be a replacement for the node 

70 if ret is not None: 

71 _invalidate_cache() 

72 node = ret 

73 if ret.__class__ != cls: 

74 # Can no longer apply the rest of the transforms. 

75 break 

76 return node 

77 

78 def _visit(self, node: nodes.NodeNG) -> SuccessfulInferenceResult: 

79 for name in node._astroid_fields: 

80 value = getattr(node, name) 

81 if TYPE_CHECKING: 

82 value = cast(_Vistables, value) 

83 

84 visited = self._visit_generic(value) 

85 if visited != value: 

86 setattr(node, name, visited) 

87 return self._transform(node) 

88 

89 @overload 

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

91 

92 @overload 

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

94 

95 @overload 

96 def _visit_generic( 

97 self, node: list[nodes.NodeNG] 

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

99 

100 @overload 

101 def _visit_generic( 

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

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

104 

105 @overload 

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

107 

108 def _visit_generic(self, node: _Vistables) -> _VisitReturns: 

109 if not node: 

110 return node 

111 if isinstance(node, list): 

112 return [self._visit_generic(child) for child in node] 

113 if isinstance(node, tuple): 

114 return tuple(self._visit_generic(child) for child in node) 

115 if isinstance(node, str): 

116 return node 

117 

118 try: 

119 return self._visit(node) 

120 except RecursionError: 

121 # Returning the node untransformed is better than giving up. 

122 warnings.warn( 

123 f"Astroid was unable to transform {node}.\n" 

124 "Some functionality will be missing unless the system recursion limit is lifted.\n" 

125 "From pylint, try: --init-hook='import sys; sys.setrecursionlimit(2000)' or higher.", 

126 UserWarning, 

127 stacklevel=0, 

128 ) 

129 return node 

130 

131 def register_transform( 

132 self, 

133 node_class: type[_SuccessfulInferenceResultT], 

134 transform: TransformFn[_SuccessfulInferenceResultT], 

135 predicate: _Predicate[_SuccessfulInferenceResultT] | None = None, 

136 ) -> None: 

137 """Register `transform(node)` function to be applied on the given node. 

138 

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

140 when called with the node as argument. 

141 

142 The transform function may return a value which is then used to 

143 substitute the original node in the tree. 

144 """ 

145 self.transforms[node_class].append((transform, predicate)) # type: ignore[index, arg-type] 

146 

147 def unregister_transform( 

148 self, 

149 node_class: type[_SuccessfulInferenceResultT], 

150 transform: TransformFn[_SuccessfulInferenceResultT], 

151 predicate: _Predicate[_SuccessfulInferenceResultT] | None = None, 

152 ) -> None: 

153 """Unregister the given transform.""" 

154 self.transforms[node_class].remove((transform, predicate)) # type: ignore[index, arg-type] 

155 

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

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

158 

159 Only the nodes which have transforms registered will actually 

160 be replaced or changed. 

161 """ 

162 return self._visit(node)