Comment Python function calls in a string with source code

Question

I'm trying to comment function calls from a Python source code.

I read the source code in a string, and then I'm trying to apply some regex to to comment method calls.

Here is an simple example working:

s = open(path, 'r').read()
# comment method calls
result = re.sub('(^[.\w]+\(.*\))', r'#\1', s, flags=re.MULTILINE)

As you can see, I'm commenting functions in the first indentation level of the source(code outside __name__ == '__main__' , methods, and classes)

But, how can I make this regex work with multiple line calls?

Example, if I have the following code into string:

Builder.load_string('''
    type: 'example'
    callback: my_callback()
''')

How can I comment each line of this call?

Answer 1

This will give you the line numbers you need to comment:

mod = "test1"
mod = importlib.import_module(mod)
p = ast.parse(inspect.getsource(mod))


for n in p.body:
    if isinstance(n, ast.Expr):
        for node in ast.walk(n):
            if isinstance(node, ast.Call):
                print(node.lineno)

For a file like:

import math

class Foo:
    def __init__(self):
        self.foo = 4
    def bar(self):
        print("hello world")

    def foo(self):
        return self.bar()

def bar():
    return 123

f = Foo()    
f.bar()    
bar()

It will output 16 and 18 the two calls.

It is just a matter of ignoring those lines and writing the new source or doing whatever you want with the updates content:

import inspect
import importlib
import ast

def get_call_lines(mod):
    mod = importlib.import_module(mod)
    p = ast.parse(inspect.getsource(mod))
    for n in p.body:
        if isinstance(n, ast.Expr):
            for node in ast.walk(n):
                if isinstance(node, ast.Call):
                    yield(node.lineno)


from StringIO import StringIO
new = StringIO()
with open("test1.py") as f:
    st = set(get_call_lines("test1"))
    for ind, line in enumerate(f, 1):
        if ind not in st:
            new.write(line)

new.seek(0)
print(new.read())

new will contain:

import math
class Foo:
    def __init__(self):
        self.foo = 4
    def bar(self):
        print("hello world")

    def foo(self):
        return self.bar()

def bar():
    return 123

f = Foo()

You could change the code at runtime with ast.NodeTransformer but it is not such a trivial task to remove nodes, the simplest approach for what you want would be to simply ignore the Call lines in the body