I have the following function
def foo():
for _ in range(1):
print("hello")
Now I want to add another print statement to print "Loop iterated" after every loop iteration. For this I define a new function that transforms foo into an ast tree, inserts the corresponding print node and then compiles the ast tree into an executable function:
def modify(func):
def wrapper():
source_code = inspect.getsource(func)
ast_tree = ast.parse(source_code)
# insert new node into ast tree
for node in ast.walk(ast_tree):
if isinstance(node, ast.For):
node.body += ast.parse("print('Loop iterated')").body
# get the compiled function
new_func = compile(ast_tree, '<ast>', 'exec')
namespace = {}
exec(new_func, globals(), namespace)
new_func = namespace[func.__name__]
return new_func()
return wrapper
This works fine as expected when using:
foo = modify(foo)
foo()
However, if I decide to use modify as a decorator:
@modify
def foo():
for _ in range(1):
print("hello")
foo()
I get the following error:
Traceback (most recent call last):
File "c:\Users\noinn\Documents\decorator_test\test.py", line 34, in <module>
foo()
File "c:\Users\noinn\Documents\decorator_test\test.py", line 25, in wrapper
return new_func()
File "c:\Users\noinn\Documents\decorator_test\test.py", line 11, in wrapper
source_code = inspect.getsource(func)
File "C:\Users\noinn\AppData\Local\Programs\Python\Python39\lib\inspect.py", line 1024, in getsource
lines, lnum = getsourcelines(object)
File "C:\Users\noinn\AppData\Local\Programs\Python\Python39\lib\inspect.py", line 1006, in getsourcelines
lines, lnum = findsource(object)
File "C:\Users\noinn\AppData\Local\Programs\Python\Python39\lib\inspect.py", line 835, in findsource
raise OSError('could not get source code')
OSError: could not get source code
Does anyone know why that error appears? Note that this does not happen If I return the original function and the error only appears once new_func() is called.
------------------------- Solution ----------------------
Simply remove the decorator from the function in the decorator itself using:
ast_tree.body[0].decorator_list = []
After some experimenting, I found out my initial hypothesis is incorrect: inspect.getsource is smart enough to retrieve the source code even if the function name is not yet set in the module globals. (surprisingly).
What happens is that the source code is retrieved along with the decorator calls as well, and when the function is called, the decorator runs again - so, it gets some re-entrancy in inspect.getsource, at which points it fails.
The solution bellow work: I just strip decorators at the module level in the retrieved source code, before feeding it to ast.parse
I also rearranged your decorator, as it would re-read the source code, and reparse the AST at each time the decorated function would be called - the way this example is, there is no need for an inner wrapper function at all. If you happen to need to parametrize your decorator, and need the inner wrapper, all the function re-writting parts, up to the creation of new_func, should be outside the wrapper, so that they run only once
import inspect
import ast
import functools
def modify(func):
func.__globals__[func.__name__] = func
source_lines = inspect.getsource(func).splitlines()
# strip decorators from the source itself
source_code = "\n".join(line for line in source_lines if not line.startswith('@'))
ast_tree = ast.parse(source_code)
# insert new node into ast tree
for node in ast.walk(ast_tree):
if isinstance(node, ast.For):
node.body += ast.parse("print('Loop iterated')").body
# get the compiled function
new_func = compile(ast_tree, '<ast>', 'exec')
namespace = {}
exec(new_func, func.__globals__, namespace)
new_func = namespace[func.__name__]
return functools.wraps(func)(new_func)
@modify
def foo():
for _ in range(1):
print("hello")
#foo = modify(foo)
foo()
initial answer
Left here for the reasoning and simpler workaround:
it can't get the source of <module>.foo function due to the simple fact that the name foo will only be defined and bound after the function definition, including all its decorators, is executed.
The name foo simply does not exist as a variable in the module at the point the decorator code is run, although it is set as the __name__ attribute in func at this point. There is nothing inspect.getsource can do about it.
In other words, decorator syntax is altogether off-limits if one is using inspect.getsource because there is nothing yet getsource can get the source of.
The good news is the workaround is simple: one have to give-up the @ decorator syntax, and apply the decorator "in the old way" (as it was before Python 2.3): one declares the function as usual, and reassign its name to the return of manually calling the decorator - this is the same as is written in your working example:
def foo(...):
...
foo = modify(foo)