I am trying to apply a decorator that changes function code, and then execute this function with changed code.
Below is the temp module with example function. I simply want the function to return [*args, *kwargs.items(), 123] instead of [*args, *kwargs.items()] if some_decorator is applied to this function.
Edit: please note that this is only a toy example, I don't intend to append new values to a list but rather rewrite a big chunk of a function.
from inspect import getsource
def some_decorator(method):
def wrapper(*args, **kwargs):
source_code = getsource(method)
code_starts_at = source_code.find('):') + 2
head = source_code[:code_starts_at]
body = source_code[code_starts_at:]
lines = body.split('\n')
return_line = [i for i in lines if 'return' in i][0]
old_expr = return_line.replace(' return ', '')
new_expr = old_expr.replace(']', ', 123]')
new_expr = head + '\n' + ' return ' + new_expr
return eval(new_expr)
return wrapper
@some_decorator
def example_func(*args, *kwargs):
return [*args, *kwargs]
A bit more of explanation: I am rewriting the original function
def example_func(*args, **kwargs):
return [*args, *kwargs.items()]
to
def example_func(*args, **kwargs):
return [*args, *kwargs.items(), 123]
I hope that eval is able to compile and run this modified function.
When I try to run it it returns a syntax error.
from temp import example_func
example_func(5)
I am aware that eval is able to cope with this:
[*args, *kwargs.items(), 123]
but only if args and kwargs are already declared. I want them to be read from example_func(args, kwargs) when I am executing example_func.
I suppose that simply writing the modified function code to a file
def example_func(*args, **kwargs):
return [*args, *kwargs.items(), 123]
and making some_decorator to execute the function with modified code instead of original one, would work just fine. However, ideally I would skip creating any intermediary files.
Is it possible to achieve that?
While you technically can do just about anything with functions and decorators in Python, you shouldn't.
In this specific case, adding an extra value to a function that returns a list is as simple as:
def some_decorator(method):
def wrapper(*args, **kwargs):
result = method(*args, **kwargs)
return result + [123]
return wrapper
This doesn't require any function code rewriting. If all you are doing is alter the inputs or the output of a function, just alter the inputs or output, and leave the function itself be.
Decorators are primarily just syntactic sugar here, a way to change
def function_name(*args, **kwargs):
# ...
function_name = create_a_wrapper_for(function_name)
into
@create_a_wrapper_for
def function_name(*args, **kwargs):
# ...
Also note that the eval() function can't alter your fuction, because eval() is strictly limited to expressions. The def syntax to create a function is a statement. Fundamentally, statements can contain expressions and other statements (e.g. if <test_expression>: <body of statements>) but expressions can't contain statements. This is why you are getting a SyntaxError exception; while [*args, *kwargs.items()] is a valid expression, return [*args, *kwargs.items()] is a statement (containing an expression):
>>> args, kwargs = (), {}
>>> eval("[*args, *kwargs.items()]")
[]
>>> eval("return [*args, *kwargs.items()]")
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
File "<string>", line 1
return [*args, *kwargs.items()]
^
SyntaxError: invalid syntax
To execute text as arbitrary Python code, you would have to use the exec() function instead, and take care to use the same namespace as the original function so any globals used in the original function still can be accessed.
For example, if the function calls another function to obtain an extra value:
def example(*args, **kwargs):
return [extra_value(), *args, *kwargs.items()]
def extra_value():
return 42
then you can’t execute the example() function in isolation; it is part of the module global namespace and looks up extra_value in that namespace when you call the function. Functions have a reference to the global namespace of the module they are created in, accessible via the function.__globals__ attribute. When you use exec() to execute a def statement creating a function, then the new function object is connected to the global namespace you passed in. Note that def creates a function object and assigns it to the function name, so you’ll have to retrieve that object from the same namespace again:
>>> namespace = {}
>>> exec("def foo(): return 42", namespace)
>>> namespace["foo"]
<function foo at 0x7f8194fb1598>
>>> namespace["foo"]()
42
>>> namespace["foo"].__globals__ is namespace
True
Next, text manipulation to rebuild Python code is very inefficient and error prone. For example, your str.replace() code would fail if the function used this instead:
def example(*args, **kwargs):
if args or kwargs:
return [
"[called with arguments:]",
*args,
*kwargs.items()
]
because now return is indented further, there are [..] brackets in a string value in the list, and the closing ] bracket of the list is on a separate line altogether.
You'd be far better off having Python compile source code into an Abstract Syntax Tree (via the ast module), then work on that tree. A directed graph of well-defined objects is much easier to manipulate than text (which is is much more flexible in how much whitespace is used, etc.). Both the above code and your example would result in a tree with a Return() node that contains an expression whose top-level would be a List() node. You can traverse that tree and find all Return() nodes and alter their List() nodes, adding an extra node to the end of the list contents.
A Python AST can be compiled to a code object (with compile()) then run through exec() (which accepts not only text, but code objects as well).
For a real-world example of a project that rewrites Python code, look at how pytest rewrites the assert statement to add extra context. They use a module import hook to do this, but as long as the source code is available for a function, you can do so with a decorator too.
Here is an example of using the ast module to alter the list in a return statement, adding in an arbitrary constant:
import ast, inspect, functools
class ReturnListInsertion(ast.NodeTransformer):
def __init__(self, value_to_insert):
self.value = value_to_insert
def visit_FunctionDef(self, node):
# remove the `some_decorator` decorator from the AST
# we don’t need to keep applying it.
if node.decorator_list:
node.decorator_list = [
n for n in node.decorator_list
if not (isinstance(n, ast.Name) and n.id == 'some_decorator')
]
return self.generic_visit(node)
def visit_Return(self, node):
if isinstance(node.value, ast.List):
# Python 3.8 and up have ast.Constant instead, which is more
# flexible.
node.value.elts.append(ast.Num(self.value))
return self.generic_visit(node)
def some_decorator(method):
source_code = inspect.getsource(method)
tree = ast.parse(source_code)
updated = ReturnListInsertion(123).visit(tree)
# fix all line number references, make it match the original
updated = ast.increment_lineno(
ast.fix_missing_locations(updated),
method.__code__.co_firstlineno
)
ast.copy_location(updated.body[0], tree)
# compile again, as a module, then execute the compiled bytecode and
# extract the new function object. Use the original namespace
# so that any global references in the function still work.
code = compile(tree, inspect.getfile(method), 'exec')
namespace = method.__globals__
exec(code, namespace)
new_function = namespace[method.__name__]
# update new function with old function attributes, name, module, documentation
# and attributes.
return functools.update_wrapper(new_function, method)
Note that this doesn't need a wrapper function. you don't need to re-work a function each time you try to call it, the decorator could do it just once and return the resulting function object directly.
Here is a demo module to try it out with:
@some_decorator
def example(*args, **kwargs):
return [extra_value(), *args, *kwargs.items()]
def extra_value():
return 42
if __name__ == '__main__':
print(example("Monty", "Python's", name="Flying circus!"))
The above outputs [42, 'Monty', "Python's", ('name', 'Flying circus!'), 123] when run.
However, it is much easier to just use the first method.
If you do want to pursue using exec() and AST manipulation, I can recommend you read up on how to do that in Green Tree Snakes.