-
Notifications
You must be signed in to change notification settings - Fork 176
/
_visitors.py
159 lines (134 loc) · 6.67 KB
/
_visitors.py
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
146
147
148
149
150
151
152
153
154
155
156
157
158
159
# 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 TYPE_CHECKING, Union
from libcst._flatten_sentinel import FlattenSentinel
from libcst._metadata_dependent import MetadataDependent
from libcst._removal_sentinel import RemovalSentinel
from libcst._typed_visitor import CSTTypedTransformerFunctions, CSTTypedVisitorFunctions
from libcst._types import CSTNodeT
if TYPE_CHECKING:
# Circular dependency for typing reasons only
from libcst._nodes.base import CSTNode # noqa: F401
CSTVisitorT = Union["CSTTransformer", "CSTVisitor"]
class CSTTransformer(CSTTypedTransformerFunctions, MetadataDependent):
"""
The low-level base visitor class for traversing a CST and creating an
updated copy of the original CST. This should be used in conjunction with
the :func:`~libcst.CSTNode.visit` method on a :class:`~libcst.CSTNode` to
visit each element in a tree starting with that node, and possibly returning
a new node in its place.
When visiting nodes using a :class:`CSTTransformer`, the return value of
:func:`~libcst.CSTNode.visit` will be a new tree with any changes made in
:func:`~libcst.CSTTransformer.on_leave` calls reflected in its children.
"""
def on_visit(self, node: "CSTNode") -> bool:
"""
Called every time a node is visited, before we've visited its children.
Returns ``True`` if children should be visited, and returns ``False``
otherwise.
"""
visit_func = getattr(self, f"visit_{type(node).__name__}", None)
if visit_func is not None:
retval = visit_func(node)
else:
retval = True
# Don't visit children IFF the visit function returned False.
return False if retval is False else True
def on_leave(
self, original_node: CSTNodeT, updated_node: CSTNodeT
) -> Union[CSTNodeT, RemovalSentinel, FlattenSentinel[CSTNodeT]]:
"""
Called every time we leave a node, after we've visited its children. If
the :func:`~libcst.CSTTransformer.on_visit` function for this node returns
``False``, this function will still be called on that node.
``original_node`` is guaranteed to be the same node as is passed to
:func:`~libcst.CSTTransformer.on_visit`, so it is safe to do state-based
checks using the ``is`` operator. Modifications should always be performed
on the ``updated_node`` so as to not overwrite changes made by child
visits.
Returning :attr:`RemovalSentinel.REMOVE` indicates that the node should be
removed from its parent. This is not always possible, and may raise an
exception if this node is required. As a convenience, you can use
:func:`RemoveFromParent` as an alias to :attr:`RemovalSentinel.REMOVE`.
"""
leave_func = getattr(self, f"leave_{type(original_node).__name__}", None)
if leave_func is not None:
updated_node = leave_func(original_node, updated_node)
return updated_node
def on_visit_attribute(self, node: "CSTNode", attribute: str) -> None:
"""
Called before a node's child attribute is visited and after we have called
:func:`~libcst.CSTTransformer.on_visit` on the node. A node's child
attributes are visited in the order that they appear in source that this
node originates from.
"""
visit_func = getattr(self, f"visit_{type(node).__name__}_{attribute}", None)
if visit_func is not None:
visit_func(node)
def on_leave_attribute(self, original_node: "CSTNode", attribute: str) -> None:
"""
Called after a node's child attribute is visited and before we have called
:func:`~libcst.CSTTransformer.on_leave` on the node.
Unlike :func:`~libcst.CSTTransformer.on_leave`, this function does
not allow modifications to the tree and is provided solely for state
management.
"""
leave_func = getattr(
self, f"leave_{type(original_node).__name__}_{attribute}", None
)
if leave_func is not None:
leave_func(original_node)
class CSTVisitor(CSTTypedVisitorFunctions, MetadataDependent):
"""
The low-level base visitor class for traversing a CST. This should be used in
conjunction with the :func:`~libcst.CSTNode.visit` method on a
:class:`~libcst.CSTNode` to visit each element in a tree starting with that
node. Unlike :class:`CSTTransformer`, instances of this class cannot modify
the tree.
When visiting nodes using a :class:`CSTVisitor`, the return value of
:func:`~libcst.CSTNode.visit` will equal the passed in tree.
"""
def on_visit(self, node: "CSTNode") -> bool:
"""
Called every time a node is visited, before we've visited its children.
Returns ``True`` if children should be visited, and returns ``False``
otherwise.
"""
visit_func = getattr(self, f"visit_{type(node).__name__}", None)
if visit_func is not None:
retval = visit_func(node)
else:
retval = True
# Don't visit children IFF the visit function returned False.
return False if retval is False else True
def on_leave(self, original_node: "CSTNode") -> None:
"""
Called every time we leave a node, after we've visited its children. If
the :func:`~libcst.CSTVisitor.on_visit` function for this node returns
``False``, this function will still be called on that node.
"""
leave_func = getattr(self, f"leave_{type(original_node).__name__}", None)
if leave_func is not None:
leave_func(original_node)
def on_visit_attribute(self, node: "CSTNode", attribute: str) -> None:
"""
Called before a node's child attribute is visited and after we have called
:func:`~libcst.CSTTransformer.on_visit` on the node. A node's child
attributes are visited in the order that they appear in source that this
node originates from.
"""
visit_func = getattr(self, f"visit_{type(node).__name__}_{attribute}", None)
if visit_func is not None:
visit_func(node)
def on_leave_attribute(self, original_node: "CSTNode", attribute: str) -> None:
"""
Called after a node's child attribute is visited and before we have called
:func:`~libcst.CSTVisitor.on_leave` on the node.
"""
leave_func = getattr(
self, f"leave_{type(original_node).__name__}_{attribute}", None
)
if leave_func is not None:
leave_func(original_node)