Search code examples
pythonabstract-syntax-tree

How to find the AST assignment node related to the instance creation?


How to select the assignment node related to the instantiation of an object?

Eg:

class A:
    
    def __init__(self, param):
        self.param = param
        curframe = inspect.currentframe().f_back
        nodes = ast.parse(inspect.getsource(curframe))
        print(nodes.body)
        
        
a = A('test')
b = 'another assign'

>>> [<_ast.ClassDef object at 0x7fdffd72f400>, <_ast.Assign object at 0x7fdffd719a60>,
>>> <_ast.Assign object at 0x7fdffd719ee0>]

I would like to select only the assignment ast object related to the instantiation of A. How to, for example, print just the specific ast.Assign object?

I tried this approach, but it fails to singularize when there are two assignment for the same class:

class A:
    
    def __init__(self, param):
        self.param = param
        curframe = inspect.currentframe().f_back
        nodes = ast.parse(inspect.getsource(curframe))
                
        for node in nodes.body:
            try:
                if node.value.func.id == self.__class__.__name__:
                    print(ast.dump(node))
            except Exception:
                pass
        
        
a = A('test')
b = 'another assign'
c = A('another test')

>>> Assign(targets=[Name(id='a', ctx=Store())], value=Call(func=Name(id='A', ctx=Load()),
>>> args=[Constant(value='test', kind=None)], keywords=[]), type_comment=None)
>>> Assign(targets=[Name(id='c', ctx=Store())], value=Call(func=Name(id='A', ctx=Load()), 
>>> args=[Constant(value='another test', kind=None)], keywords=[]), type_comment=None)
>>> Assign(targets=[Name(id='a', ctx=Store())], value=Call(func=Name(id='A', ctx=Load()), 
>>> args=[Constant(value='test', kind=None)], keywords=[]), type_comment=None)
>>> Assign(targets=[Name(id='c', ctx=Store())], value=Call(func=Name(id='A', ctx=Load()), 
>>> args=[Constant(value='another test', kind=None)], keywords=[]), type_comment=None)

I think there is another easy way to achive this. Does anyone know a better approach?

Thanks in advance.


Solution

  • Since you have two instantiations of objects of class A, and the code makes no differentiation to the actual line being executed, both assignment statements are printed twice.

    Changing the code a little bit by retaining self.nodes (easy for debugging) and adding the current line within the outer frame, we can make sure that the actual assignment being executed is printed only.

    import ast
    import inspect
    
    class A:
        def __init__(self, param):
            self.param = param
            curframe = inspect.currentframe().f_back
            curline = inspect.currentframe().f_back.f_lineno
            self.nodes = ast.parse(inspect.getsource(curframe))
    
            for node in self.nodes.body:
                if node.lineno == curline:
                    print(ast.dump(node))
    
    a = A('test')
    b = 'another assign'
    c = A('another test')
    

    Output (reformatted for presentation):

    # Assign(targets=[Name(id='a', ctx=Store())],
             value=Call(func=Name(id='A', ctx=Load()),
             args=[Constant(value='test')], keywords=[]))
    # Assign(targets=[Name(id='c', ctx=Store())],
             value=Call(func=Name(id='A', ctx=Load()),
             args=[Constant(value='another test')], keywords=[]))