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

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 

7from collections import defaultdict 

8from collections.abc import Callable 

9from typing import TYPE_CHECKING, List, Optional, Tuple, TypeVar, Union, cast, overload 

10 

11from astroid.context import _invalidate_cache 

12from astroid.typing import SuccessfulInferenceResult, TransformFn 

13 

14if TYPE_CHECKING: 

15 from astroid import nodes 

16 

17 _SuccessfulInferenceResultT = TypeVar( 

18 "_SuccessfulInferenceResultT", bound=SuccessfulInferenceResult 

19 ) 

20 _Predicate = Optional[Callable[[_SuccessfulInferenceResultT], bool]] 

21 

22_Vistables = Union[ 

23 "nodes.NodeNG", List["nodes.NodeNG"], Tuple["nodes.NodeNG", ...], str, None 

24] 

25_VisitReturns = Union[ 

26 SuccessfulInferenceResult, 

27 List[SuccessfulInferenceResult], 

28 Tuple[SuccessfulInferenceResult, ...], 

29 str, 

30 None, 

31] 

32 

33 

34class TransformVisitor: 

35 """A visitor for handling transforms. 

36 

37 The standard approach of using it is to call 

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

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

40 transforms for each encountered node. 

41 

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

43 """ 

44 

45 def __init__(self) -> None: 

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

47 # Refer to register_transform and unregister_transform for the correct types 

48 self.transforms: defaultdict[ 

49 type[SuccessfulInferenceResult], 

50 list[ 

51 tuple[ 

52 TransformFn[SuccessfulInferenceResult], 

53 _Predicate[SuccessfulInferenceResult], 

54 ] 

55 ], 

56 ] = defaultdict(list) 

57 

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

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

60 transformed node. 

61 """ 

62 cls = node.__class__ 

63 

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

65 if predicate is None or predicate(node): 

66 ret = transform_func(node) 

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

68 # expected to be a replacement for the node 

69 if ret is not None: 

70 _invalidate_cache() 

71 node = ret 

72 if ret.__class__ != cls: 

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

74 break 

75 return node 

76 

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

78 for name in node._astroid_fields: 

79 value = getattr(node, name) 

80 value = cast(_Vistables, value) 

81 visited = self._visit_generic(value) 

82 if visited != value: 

83 setattr(node, name, visited) 

84 return self._transform(node) 

85 

86 @overload 

87 def _visit_generic(self, node: None) -> None: 

88 ... 

89 

90 @overload 

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

92 ... 

93 

94 @overload 

95 def _visit_generic( 

96 self, node: list[nodes.NodeNG] 

97 ) -> list[SuccessfulInferenceResult]: 

98 ... 

99 

100 @overload 

101 def _visit_generic( 

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

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

104 ... 

105 

106 @overload 

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

108 ... 

109 

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

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 not node or isinstance(node, str): 

116 return node 

117 

118 return self._visit(node) 

119 

120 def register_transform( 

121 self, 

122 node_class: type[_SuccessfulInferenceResultT], 

123 transform: TransformFn[_SuccessfulInferenceResultT], 

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

125 ) -> None: 

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

127 

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

129 when called with the node as argument. 

130 

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

132 substitute the original node in the tree. 

133 """ 

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

135 

136 def unregister_transform( 

137 self, 

138 node_class: type[_SuccessfulInferenceResultT], 

139 transform: TransformFn[_SuccessfulInferenceResultT], 

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

141 ) -> None: 

142 """Unregister the given transform.""" 

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

144 

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

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

147 

148 Only the nodes which have transforms registered will actually 

149 be replaced or changed. 

150 """ 

151 return self._visit(node)