Search code examples
pythonabstract-syntax-tree

Assignment operator overloading in python Abstract Syntax Trees


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")


Solution

  • 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')