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

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, List, Optional, Tuple, 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 value = cast(_Vistables, value) 

82 visited = self._visit_generic(value) 

83 if visited != value: 

84 setattr(node, name, visited) 

85 return self._transform(node) 

86 

87 @overload 

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

89 

90 @overload 

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

92 

93 @overload 

94 def _visit_generic( 

95 self, node: list[nodes.NodeNG] 

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

97 

98 @overload 

99 def _visit_generic( 

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

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

102 

103 @overload 

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

105 

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

107 if isinstance(node, list): 

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

109 if isinstance(node, tuple): 

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

111 if not node or isinstance(node, str): 

112 return node 

113 

114 try: 

115 return self._visit(node) 

116 except RecursionError: 

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

118 warnings.warn( 

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

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

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

122 UserWarning, 

123 stacklevel=0, 

124 ) 

125 return node 

126 

127 def register_transform( 

128 self, 

129 node_class: type[_SuccessfulInferenceResultT], 

130 transform: TransformFn[_SuccessfulInferenceResultT], 

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

132 ) -> None: 

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

134 

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

136 when called with the node as argument. 

137 

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

139 substitute the original node in the tree. 

140 """ 

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

142 

143 def unregister_transform( 

144 self, 

145 node_class: type[_SuccessfulInferenceResultT], 

146 transform: TransformFn[_SuccessfulInferenceResultT], 

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

148 ) -> None: 

149 """Unregister the given transform.""" 

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

151 

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

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

154 

155 Only the nodes which have transforms registered will actually 

156 be replaced or changed. 

157 """ 

158 return self._visit(node)