Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/libcst/_batched_visitor.py: 38%

1# Copyright (c) Meta Platforms, Inc. and affiliates. 

2# 

3# This source code is licensed under the MIT license found in the 

4# LICENSE file in the root directory of this source tree. 

5 

6import inspect 

7from typing import ( 

8 Callable, 

9 cast, 

10 Iterable, 

11 List, 

12 Mapping, 

13 MutableMapping, 

14 Optional, 

15 TYPE_CHECKING, 

16) 

17 

18from libcst._metadata_dependent import MetadataDependent 

19from libcst._typed_visitor import CSTTypedVisitorFunctions 

20from libcst._visitors import CSTNodeT, CSTVisitor 

21 

22if TYPE_CHECKING: 

23 from libcst._nodes.base import CSTNode # noqa: F401 

24 

25VisitorMethod = Callable[["CSTNode"], None] 

26_VisitorMethodCollection = Mapping[str, List[VisitorMethod]] 

27 

28 

29class BatchableCSTVisitor(CSTTypedVisitorFunctions, MetadataDependent): 

30 """ 

31 The low-level base visitor class for traversing a CST as part of a batched 

32 set of traversals. This should be used in conjunction with the 

33 :func:`~libcst.visit_batched` function or the 

34 :func:`~libcst.MetadataWrapper.visit_batched` method from 

35 :class:`~libcst.MetadataWrapper` to visit a tree. 

36 Instances of this class cannot modify the tree. 

37 """ 

38 

39 def get_visitors(self) -> Mapping[str, VisitorMethod]: 

40 """ 

41 Returns a mapping of all the ``visit_<Type[CSTNode]>``, 

42 ``visit_<Type[CSTNode]>_<attribute>``, ``leave_<Type[CSTNode]>`` and 

43 `leave_<Type[CSTNode]>_<attribute>`` methods defined by this visitor, 

44 excluding all empty stubs. 

45 """ 

46 

47 methods = inspect.getmembers( 

48 self, 

49 lambda m: ( 

50 inspect.ismethod(m) 

51 and (m.__name__.startswith("visit_") or m.__name__.startswith("leave_")) 

52 and not getattr(m, "_is_no_op", False) 

53 ), 

54 ) 

55 

56 # TODO: verify all visitor methods reference valid node classes. 

57 # for name, __ in methods: 

58 # ... 

59 

60 return dict(methods) 

61 

62 

63def visit_batched( 

64 node: CSTNodeT, 

65 batchable_visitors: Iterable[BatchableCSTVisitor], 

66 before_visit: Optional[VisitorMethod] = None, 

67 after_leave: Optional[VisitorMethod] = None, 

68) -> CSTNodeT: 

69 """ 

70 Do a batched traversal over ``node`` with all ``visitors``. 

71 

72 ``before_visit`` and ``after_leave`` are provided as optional hooks to 

73 execute before the ``visit_<Type[CSTNode]>`` and after the 

74 ``leave_<Type[CSTNode]>`` methods from each visitor in ``visitor`` are 

75 executed by the batched visitor. 

76 

77 This function does not handle metadata dependency resolution for ``visitors``. 

78 See :func:`~libcst.MetadataWrapper.visit_batched` from 

79 :class:`~libcst.MetadataWrapper` for batched traversal with metadata dependency 

80 resolution. 

81 """ 

82 visitor_methods = _get_visitor_methods(batchable_visitors) 

83 batched_visitor = _BatchedCSTVisitor( 

84 visitor_methods, before_visit=before_visit, after_leave=after_leave 

85 ) 

86 return cast(CSTNodeT, node.visit(batched_visitor)) 

87 

88 

89def _get_visitor_methods( 

90 batchable_visitors: Iterable[BatchableCSTVisitor], 

91) -> _VisitorMethodCollection: 

92 """ 

93 Gather all ``visit_<Type[CSTNode]>``, ``visit_<Type[CSTNode]>_<attribute>``, 

94 ``leave_<Type[CSTNode]>`` amd `leave_<Type[CSTNode]>_<attribute>`` methods 

95 from ``batchabled_visitors``. 

96 """ 

97 visitor_methods: MutableMapping[str, List[VisitorMethod]] = {} 

98 for bv in batchable_visitors: 

99 for name, fn in bv.get_visitors().items(): 

100 visitor_methods.setdefault(name, []).append(fn) 

101 return visitor_methods 

102 

103 

104class _BatchedCSTVisitor(CSTVisitor): 

105 """ 

106 Internal visitor class to perform batched traversal over a tree. 

107 """ 

108 

109 visitor_methods: _VisitorMethodCollection 

110 before_visit: Optional[VisitorMethod] 

111 after_leave: Optional[VisitorMethod] 

112 

113 def __init__( 

114 self, 

115 visitor_methods: _VisitorMethodCollection, 

116 *, 

117 before_visit: Optional[VisitorMethod] = None, 

118 after_leave: Optional[VisitorMethod] = None, 

119 ) -> None: 

120 super().__init__() 

121 self.visitor_methods = visitor_methods 

122 self.before_visit = before_visit 

123 self.after_leave = after_leave 

124 

125 def on_visit(self, node: "CSTNode") -> bool: 

126 """ 

127 Call appropriate visit methods on node before visiting children. 

128 """ 

129 before_visit = self.before_visit 

130 if before_visit is not None: 

131 before_visit(node) 

132 type_name = type(node).__name__ 

133 for v in self.visitor_methods.get(f"visit_{type_name}", []): 

134 v(node) 

135 return True 

136 

137 def on_leave(self, original_node: "CSTNode") -> None: 

138 """ 

139 Call appropriate leave methods on node after visiting children. 

140 """ 

141 type_name = type(original_node).__name__ 

142 for v in self.visitor_methods.get(f"leave_{type_name}", []): 

143 v(original_node) 

144 after_leave = self.after_leave 

145 if after_leave is not None: 

146 after_leave(original_node) 

147 

148 def on_visit_attribute(self, node: "CSTNode", attribute: str) -> None: 

149 """ 

150 Call appropriate visit attribute methods on node before visiting 

151 attribute's children. 

152 """ 

153 type_name = type(node).__name__ 

154 for v in self.visitor_methods.get(f"visit_{type_name}_{attribute}", []): 

155 v(node) 

156 

157 def on_leave_attribute(self, original_node: "CSTNode", attribute: str) -> None: 

158 """ 

159 Call appropriate leave attribute methods on node after visiting 

160 attribute's children. 

161 """ 

162 type_name = type(original_node).__name__ 

163 for v in self.visitor_methods.get(f"leave_{type_name}_{attribute}", []): 

164 v(original_node)