I want to overload assignment operator in python on the fly using Abstract Syntax Trees
import ast
import astunparse
class OverloadAssignments(ast.NodeTransformer):
def visit_Assign(self, node):
if isinstance(node, ast.Assign) and node.targets:
funcs = node.targets[0]
slot_name_candidate = astunparse.unparse(funcs).strip()
if isinstance(funcs, ast.Name) and "_slot" in slot_name_candidate:
slot_name = ast.Constant(value=slot_name_candidate)
context_variable = ast.Constant(value=astunparse.unparse(node.value).strip())
return ast.Expr([ast.Call(func=ast.Name(id='copy_variable_value', ctx=ast.Load),
args=[slot_name, context_variable], keywords=[])])
else:
return node
return node
assignment_overloader = OverloadAssignments()
code_chunk = "town_slot=cxt.my_town"
tree = ast.parse(code_chunk)
tree = assignment_overloader.visit(tree)
I use parseprint
function for pretty printing code tree structure from here
https://bitbucket.org/takluyver/greentreesnakes/src/master/astpp.py
http://alexleone.blogspot.co.uk/2010/01/python-ast-pretty-printer.html
which gives me the result
parseprint(tree)
Module(body=[
Expr(value=[
Call(func=Name(id='copy_variable_value', ctx=<class 'ast.Load'>), args=[
Constant(value='town_slot', kind=None),
Constant(value='cxt.my_town', kind=None),
], keywords=[]),
]),
], type_ignores=[])
Than I need to unparse code to string. I do it with another python package:
astunparse.unparse(tree)
AttributeError: 'Unparser' object has no attribute '_str'
which fails.
What does cause astunparse to fail in this case?
How do I correctly unparse the above code?
I expect astunparse
to produce the following code chunk:
copy_variable_value("town_slot", "cxt.my_town")
You do not need to use astunparse
, the ast
module includes an unparse
method:
import ast
class AssignOverload(ast.NodeTransformer):
def visit_Assign(self, node):
return ast.Call(func=ast.Name(id='copy_variable_value'),
args=[ast.Constant(value=ast.unparse(i)) for i in [*node.targets, node.value]],
keywords=[])
code_chunk = "town_slot=cxt.my_town"
a = AssignOverload()
result = a.visit(ast.parse(code_chunk))
print(ast.unparse(result))
Output:
copy_variable_value('town_slot', 'cxt.my_town')