Search code examples
python-3.xcode-generation

Runtime compiling a function with arguments in python


I am trying to use compile to runtime generate a Python function accepting arguments as follows.

import types
import ast

code = compile("def add(a, b): return a + b", '<string>', 'exec')
fn = types.FunctionType(code, {}, name="add")
print(fn(4, 2))

But it fails with

TypeError: <module>() takes 0 positional arguments but 2 were given

Is there anyway to compile a function accepting arguments using this way or is there any other way to do that?


Solution

  • Compile returns the code object to create a module. In Python 3.6, if you were to disassemble your code object:

    >>> import dis
    >>> dis.dis(fn)
     0 LOAD_CONST    0 (<code object add at ...., file "<string>" ...>)
     2 LOAD_CONST    1 ('add')
     4 MAKE_FUNCTION 0
     6 STORE_NAME    0 (add)
     8 LOAD_CONST    2 (None)
    10 RETURN_VALUE
    

    That literally translates to make function; name it 'add'; return None.

    This code means that your function runs the creation of the module, not returning a module or function itself. So essentially, what you're actually doing is equivalent to the following:

    def f():
        def add(a, b):
            return a + b
    
    print(f(4, 2))
    

    For the question of how do you work around, the answer is it depends on what you want to do. For instance, if you want to compile a function using compile, the simple answer is you won't be able to without doing something similar to the following.

    # 'code' is the result of the call to compile.
    # In this case we know it is the first constant (from dis),
    # so we will go and extract it's value
    f_code = code.co_consts[0]
    add = FunctionType(f_code, {}, "add")
    
    >>> add(4, 2)
    6
    

    Since defining a function in Python requires running Python code (there is no static compilation by default other than compiling to bytecode), you can pass in custom globals and locals dictionaries, and then extract the values from those.

    glob, loc = {}, {}
    exec(code, glob, loc)
    
    >>> loc['add'](4, 2)
    6
    

    But the real answer is if you want to do this, the simplest way is generally to generate Abstract Syntax Trees using the ast module, and compiling that into module code and evaluating or executing the module.

    If you want to do bytecode transformation, I'd suggest looking at the codetransformer package on PyPi.

    TL;DR using compile will only ever return code for a module, and most serious code generation is done either with ASTs or by manipulating byte codes.