Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/libcst/codemod/visitors/_gather_exports.py: 26%

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# 

6from typing import Set, Union 

7 

8import libcst as cst 

9import libcst.matchers as m 

10from libcst.codemod._context import CodemodContext 

11from libcst.codemod._visitor import ContextAwareVisitor 

12from libcst.helpers import get_full_name_for_node 

13 

14 

15class GatherExportsVisitor(ContextAwareVisitor): 

16 """ 

17 Gathers all explicit exports in a module and stores them as attributes on the 

18 instance. Intended to be instantiated and passed to a :class:`~libcst.Module` 

19 :meth:`~libcst.CSTNode.visit` method in order to gather up information about 

20 exports specified in an ``__all__`` variable inside a module. 

21 

22 After visiting a module the following attributes will be populated: 

23 

24 explicit_exported_objects 

25 A sequence of strings representing objects that the module exports 

26 directly. Note that when ``__all__`` is absent, this attribute does not 

27 store default exported objects by name. 

28 

29 For more information on ``__all__``, please see Python's `Modules Documentation 

30 <https://docs.python.org/3/tutorial/modules.html>`_. 

31 """ 

32 

33 def __init__(self, context: CodemodContext) -> None: 

34 super().__init__(context) 

35 # Track any re-exported objects in an __all__ reference and whether 

36 # they're defined or not 

37 self.explicit_exported_objects: Set[str] = set() 

38 

39 # Presumably at some point in the future it would be useful to grab 

40 # a list of all implicitly exported objects. That would go here as 

41 # well and would follow Python's rule for importing objects that 

42 # do not start with an underscore. Because of that, I named the above 

43 # `explicit_exported_objects` instead of just `exported_objects` so 

44 # that we have a reasonable place to put implicit objects in the future. 

45 

46 # Internal bookkeeping 

47 self._is_assigned_export: Set[Union[cst.Tuple, cst.List, cst.Set]] = set() 

48 self._in_assigned_export: Set[Union[cst.Tuple, cst.List, cst.Set]] = set() 

49 

50 def visit_AnnAssign(self, node: cst.AnnAssign) -> bool: 

51 value = node.value 

52 if value: 

53 if self._handle_assign_target(node.target, value): 

54 return True 

55 return False 

56 

57 def visit_AugAssign(self, node: cst.AugAssign) -> bool: 

58 if m.matches( 

59 node, 

60 m.AugAssign( 

61 target=m.Name("__all__"), 

62 operator=m.AddAssign(), 

63 value=m.List() | m.Tuple(), 

64 ), 

65 ): 

66 value = node.value 

67 if isinstance(value, (cst.List, cst.Tuple)): 

68 self._is_assigned_export.add(value) 

69 return True 

70 return False 

71 

72 def visit_Assign(self, node: cst.Assign) -> bool: 

73 for target_node in node.targets: 

74 if self._handle_assign_target(target_node.target, node.value): 

75 return True 

76 return False 

77 

78 def _handle_assign_target( 

79 self, target: cst.BaseExpression, value: cst.BaseExpression 

80 ) -> bool: 

81 target_name = get_full_name_for_node(target) 

82 if target_name == "__all__": 

83 # Assignments such as `__all__ = ["os"]` 

84 # or `__all__ = exports = ["os"]` 

85 if isinstance(value, (cst.List, cst.Tuple, cst.Set)): 

86 self._is_assigned_export.add(value) 

87 return True 

88 elif isinstance(target, cst.Tuple) and isinstance(value, cst.Tuple): 

89 # Assignments such as `__all__, x = ["os"], []` 

90 for element_idx, element_node in enumerate(target.elements): 

91 element_name = get_full_name_for_node(element_node.value) 

92 if element_name == "__all__": 

93 element_value = value.elements[element_idx].value 

94 if isinstance(element_value, (cst.List, cst.Tuple, cst.Set)): 

95 self._is_assigned_export.add(value) 

96 self._is_assigned_export.add(element_value) 

97 return True 

98 return False 

99 

100 def visit_List(self, node: cst.List) -> bool: 

101 if node in self._is_assigned_export: 

102 self._in_assigned_export.add(node) 

103 return True 

104 return False 

105 

106 def leave_List(self, original_node: cst.List) -> None: 

107 self._is_assigned_export.discard(original_node) 

108 self._in_assigned_export.discard(original_node) 

109 

110 def visit_Tuple(self, node: cst.Tuple) -> bool: 

111 if node in self._is_assigned_export: 

112 self._in_assigned_export.add(node) 

113 return True 

114 return False 

115 

116 def leave_Tuple(self, original_node: cst.Tuple) -> None: 

117 self._is_assigned_export.discard(original_node) 

118 self._in_assigned_export.discard(original_node) 

119 

120 def visit_Set(self, node: cst.Set) -> bool: 

121 if node in self._is_assigned_export: 

122 self._in_assigned_export.add(node) 

123 return True 

124 return False 

125 

126 def leave_Set(self, original_node: cst.Set) -> None: 

127 self._is_assigned_export.discard(original_node) 

128 self._in_assigned_export.discard(original_node) 

129 

130 def visit_SimpleString(self, node: cst.SimpleString) -> bool: 

131 self._handle_string_export(node) 

132 return False 

133 

134 def visit_ConcatenatedString(self, node: cst.ConcatenatedString) -> bool: 

135 self._handle_string_export(node) 

136 return False 

137 

138 def _handle_string_export( 

139 self, node: Union[cst.SimpleString, cst.ConcatenatedString] 

140 ) -> None: 

141 if self._in_assigned_export: 

142 name = node.evaluated_value 

143 if not isinstance(name, str): 

144 return 

145 self.explicit_exported_objects.add(name)