def test_delete_print_statements(self):
def is_print_statement(ast: AST) -> bool:
"""Return TRUE if AST is an statement calling the print function."""
if isinstance(ast, ExpressionStatementAST):
fn_calls = [
c.call_function().source_text for c in ast.call_asts()
]
return "print" in fn_calls
return False
def delete_print_statements(ast: AST) -> Optional[LiteralOrAST]:
"""Delete all print statements from the children of AST."""
if isinstance(ast, RootAST) or isinstance(ast, CompoundAST):
# Build a list of new children under the AST, eliding print statements.
new_children = [
c for c in ast.children if not is_print_statement(c)
]
# Special case; if no children remain, add a "pass" statement nop to
# avoid syntax errors.
new_children = new_children if new_children else ["pass\n"]
return AST.copy(ast, children=new_children)
transformed = AST.transform(self.root, delete_print_statements)
expected = slurp(DATA_DIR / "transform" / "delete_print_statements.py")
self.assertEqual(transformed.source_text, expected)
def test_remove_cast(self):
def remove_cast(ast: AST):
if isinstance(ast, CCastExpression):
return ast.value
transformed = AST.transform(self.root, remove_cast)
expected = slurp(DATA_DIR / "transform" / "ccast-transformed.c")
self.assertEqual(transformed.source_text, expected)
def test_transform_x_to_y(self):
def x_to_y(ast: AST) -> Optional[LiteralOrAST]:
"""Convert 'x' identifier ASTs to 'y'."""
if isinstance(ast, IdentifierAST) and "x" == ast.source_text:
return "y"
transformed = AST.transform(self.root, x_to_y)
expected = slurp(DATA_DIR / "transform" / "transform_x_to_y.py")
self.assertEqual(transformed.source_text, expected)
def test_transform_x_to_z_gt(self):
def x_to_z_gt(ast: AST) -> Optional[LiteralOrAST]:
"""Convert 'x' identifiers in greater than operations to 'z'."""
if isinstance(ast, PythonComparisonOperator):
lhs, *rest = ast.child_slot("CHILDREN")
operator, *_ = ast.child_slot("OPERATORS")
if isinstance(operator,
PythonGreaterThan) and lhs.source_text == "x":
return AST.copy(ast, children=["z", *rest])
transformed = AST.transform(self.root, x_to_z_gt)
expected = slurp(DATA_DIR / "transform" / "transform_x_to_z_gt.py")
self.assertEqual(transformed.source_text, expected)