def on_leave(
self, original_node: CSTNodeT, updated_node: CSTNodeT
) -> Union[cst.Import, cst.ImportFrom, CSTNodeT, RemovalSentinel]:
if isinstance(updated_node, cst.Import):
for alias in updated_node.names:
name = alias.name
if isinstance(name, cst.Name) and name.value == "b":
return updated_node.remove()
elif isinstance(updated_node, cst.ImportFrom):
module = updated_node.module
if isinstance(module, cst.Name) and module.value == "e":
return updated_node.remove()
return updated_node
def validate_node(
self,
node: CSTNodeT,
code: str,
parser: Optional[Callable[[str], CSTNodeT]] = None,
expected_position: Optional[CodeRange] = None,
) -> None:
node.validate_types_deep()
self.__assert_codegen(node, code, expected_position)
if parser is not None:
parsed_node = parser(code)
self.assertEqual(parsed_node, node)
# Tests of children should unwrap DummyIndentedBlock first, because we don't
# want to test DummyIndentedBlock's behavior.
unwrapped_node = node
while isinstance(unwrapped_node, DummyIndentedBlock):
unwrapped_node = unwrapped_node.child
self.__assert_children_match_codegen(unwrapped_node)
self.__assert_children_match_fields(unwrapped_node)
self.__assert_visit_returns_identity(unwrapped_node)
def on_leave(self, original_node: CSTNodeT,
updated_node: CSTNodeT) -> Union[CSTNodeT, RemovalSentinel]:
self.test.assertTrue(original_node.deep_equals(updated_node))
# Don't allow type checkers to accidentally refine our return type.
return_node = updated_node
if isinstance(updated_node, cst.Pass):
self.assert_counter(3)
elif isinstance(updated_node, cst.Newline):
self.assert_counter(5)
elif isinstance(updated_node, cst.SimpleStatementLine):
self.assert_counter(6)
elif isinstance(updated_node, cst.Module):
self.assert_counter(7)
return return_node
def visit_required(parent: "CSTNode", fieldname: str, node: CSTNodeT,
visitor: "CSTVisitorT") -> CSTNodeT:
"""
Given a node, visits the node using `visitor`. If removal is attempted by the
visitor, an exception is raised.
"""
visitor.on_visit_attribute(parent, fieldname)
result = node.visit(visitor)
if isinstance(result, RemovalSentinel):
raise TypeError(
f"We got a RemovalSentinel while visiting a {type(node).__name__}. This "
+ "node's parent does not allow it to be removed.")
visitor.on_leave_attribute(parent, fieldname)
return result
def on_leave(self, original_node: CSTNodeT,
updated_node: CSTNodeT) -> Union[CSTNodeT, RemovalSentinel]:
if isinstance(updated_node, cst.Continue):
return updated_node.remove()
else:
return updated_node