1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
|
# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
#
from typing import Set, Union
import libcst as cst
import libcst.matchers as m
from libcst.codemod._context import CodemodContext
from libcst.codemod._visitor import ContextAwareVisitor
from libcst.helpers import get_full_name_for_node
class GatherExportsVisitor(ContextAwareVisitor):
"""
Gathers all explicit exports in a module and stores them as attributes on the
instance. Intended to be instantiated and passed to a :class:`~libcst.Module`
:meth:`~libcst.CSTNode.visit` method in order to gather up information about
exports specified in an ``__all__`` variable inside a module.
After visiting a module the following attributes will be populated:
explicit_exported_objects
A sequence of strings representing objects that the module exports
directly. Note that when ``__all__`` is absent, this attribute does not
store default exported objects by name.
For more information on ``__all__``, please see Python's `Modules Documentation
<https://docs.python.org/3/tutorial/modules.html>`_.
"""
def __init__(self, context: CodemodContext) -> None:
super().__init__(context)
# Track any re-exported objects in an __all__ reference and whether
# they're defined or not
self.explicit_exported_objects: Set[str] = set()
# Presumably at some point in the future it would be useful to grab
# a list of all implicitly exported objects. That would go here as
# well and would follow Python's rule for importing objects that
# do not start with an underscore. Because of that, I named the above
# `explicit_exported_objects` instead of just `exported_objects` so
# that we have a reasonable place to put implicit objects in the future.
# Internal bookkeeping
self._is_assigned_export: Set[Union[cst.Tuple, cst.List, cst.Set]] = set()
self._in_assigned_export: Set[Union[cst.Tuple, cst.List, cst.Set]] = set()
def visit_AnnAssign(self, node: cst.AnnAssign) -> bool:
value = node.value
if value:
if self._handle_assign_target(node.target, value):
return True
return False
def visit_AugAssign(self, node: cst.AugAssign) -> bool:
if m.matches(
node,
m.AugAssign(
target=m.Name("__all__"),
operator=m.AddAssign(),
value=m.List() | m.Tuple(),
),
):
value = node.value
if isinstance(value, (cst.List, cst.Tuple)):
self._is_assigned_export.add(value)
return True
return False
def visit_Assign(self, node: cst.Assign) -> bool:
for target_node in node.targets:
if self._handle_assign_target(target_node.target, node.value):
return True
return False
def _handle_assign_target(
self, target: cst.BaseExpression, value: cst.BaseExpression
) -> bool:
target_name = get_full_name_for_node(target)
if target_name == "__all__":
# Assignments such as `__all__ = ["os"]`
# or `__all__ = exports = ["os"]`
if isinstance(value, (cst.List, cst.Tuple, cst.Set)):
self._is_assigned_export.add(value)
return True
elif isinstance(target, cst.Tuple) and isinstance(value, cst.Tuple):
# Assignments such as `__all__, x = ["os"], []`
for element_idx, element_node in enumerate(target.elements):
element_name = get_full_name_for_node(element_node.value)
if element_name == "__all__":
element_value = value.elements[element_idx].value
if isinstance(element_value, (cst.List, cst.Tuple, cst.Set)):
self._is_assigned_export.add(value)
self._is_assigned_export.add(element_value)
return True
return False
def visit_List(self, node: cst.List) -> bool:
if node in self._is_assigned_export:
self._in_assigned_export.add(node)
return True
return False
def leave_List(self, original_node: cst.List) -> None:
self._is_assigned_export.discard(original_node)
self._in_assigned_export.discard(original_node)
def visit_Tuple(self, node: cst.Tuple) -> bool:
if node in self._is_assigned_export:
self._in_assigned_export.add(node)
return True
return False
def leave_Tuple(self, original_node: cst.Tuple) -> None:
self._is_assigned_export.discard(original_node)
self._in_assigned_export.discard(original_node)
def visit_Set(self, node: cst.Set) -> bool:
if node in self._is_assigned_export:
self._in_assigned_export.add(node)
return True
return False
def leave_Set(self, original_node: cst.Set) -> None:
self._is_assigned_export.discard(original_node)
self._in_assigned_export.discard(original_node)
def visit_SimpleString(self, node: cst.SimpleString) -> bool:
self._handle_string_export(node)
return False
def visit_ConcatenatedString(self, node: cst.ConcatenatedString) -> bool:
self._handle_string_export(node)
return False
def _handle_string_export(
self, node: Union[cst.SimpleString, cst.ConcatenatedString]
) -> None:
if self._in_assigned_export:
name = node.evaluated_value
if not isinstance(name, str):
return
self.explicit_exported_objects.add(name)
|
