Partial evaluation of Python code#

Introduction#

This library allows one to perform code specialization at run-time, turning this:

@inline
def power(x, n):
    if n == 0:
        return 1
    elif n % 2 == 0:
        v = power(x, n // 2)
        return v * v
    else:
        return x * power(x, n - 1)

with n set to, say, 27, into this:

def power_27(x):
    __peval_return_7 = (x * 1)
    __peval_return_6 = (__peval_return_7 * __peval_return_7)
    __peval_return_5 = (x * __peval_return_6)
    __peval_return_4 = (__peval_return_5 * __peval_return_5)
    __peval_return_3 = (__peval_return_4 * __peval_return_4)
    __peval_return_1 = (x * __peval_return_3)
    __peval_return_2 = (__peval_return_1 * __peval_return_1)
    return (x * __peval_return_2)

The resulting code runs 6 times faster under CPython 3.5.1.

The API is identical to that of functools.partial():

import peval
power_27 = peval.partial_apply(power, n=27)

You must mark the functions that you want inlined (maybe recursively) with peval.inline(). If you want some function to be evaluated during partial evaluation, mark it with peval.pure() (if it is not, in fact, pure, the results are unpredictable).

Implementation details#

The partial evaluation is performed on an AST using a range of optimizations:

constant propagation
constant folding
unreachable code elimination
function inlining

Restrictions on functions#

Currently not all functions can be partially evaluated or inlined. Hopefully, these restrictions will be lifted in the future as the package improves.

The inlined functions:

cannot be async functions;
cannot be generators;
cannot have nested definitions (lambdas, functions or classes);
cannot have closures;
cannot have decorators and annotations (they will be ignored).

The partially evaluated functions:

cannot be async functions;
cannot have nested definitions (lambdas, functions or classes);
cannot have decorators and annotations (they will be ignored).

Also note that the partial evaluated function loses the connection to the globals dictionary of the original function. Therefore, it cannot reassign global variables (the copy is shallow, so mutation of global objects is still possible).

API reference#

High-level API#

peval.partial_eval(func: Callable) → Callable#: Returns a partially evaluated version of func, using the values of associated global and closure variables.

peval.partial_apply(func: Callable, *args, **kwds) → Callable#: Same as partial_eval(), but in addition uses the provided values of positional and keyword arguments in the partial evaluation.

Tags#

peval.inline(func: Callable) → Callable#: Marks the function for inlining.

peval.pure(func: Callable) → Callable#: Marks the function as pure (not having any side effects, except maybe argument mutation).

Helper functions#

peval.getsource(func: Callable) → str#: Returns the source of a function func. Falls back to inspect.getsource() for regular functions, but can also return the source of a partially evaluated function.

peval.specialize_on(names: Union[str, Tuple[str, str]], maxsize=None) → Callable#: A decorator that wraps a function, partially evaluating it with the parameters defined by names (can be a string or an iterable of strings) being fixed. The partially evaluated versions are cached based on the values of these parameters using functools.lru_cache with the provided maxsize (consequently, these values should be hashable).

Low-level tools#

peval.ast_walker(handler)#

A generic AST walker decorator. Decorates either a function or a class (if dispatching based on node type is required). handler will be wrapped in a Dispatcher instance; see Dispatcher for the details of the required class structure.

Returns a callable with the signature:

def walker(state, node, ctx=None)

Parameters:

state – a dictionary with the state which will be passed to every handler call. It will be converted into a immutableadict object at the start of the traversal. Handlers can update it by returning a modified version.
node – an ast.AST object to traverse.
ctx – a dictionary with the global context which will be passed to every handler call. It will be converted into a immutableadict object at the start of the traversal.

Returns:

a tuple (state, new_node), where state is the same object which was passed as the corresponding parameter. Does not mutate node.

handler will be invoked for every node during the AST traversal (depth-first, pre-order). The handler function, if it is a function, or its static methods, if it is a class must have the signature:

def handler([state, node, ctx, prepend, visit_after, visiting_after,
    skip_fields, walk_field,] **kwds)

The names of the arguments must be exactly as written here, but their order is not significant (they will be passed as keywords).

If handler is a class, the default handler is a “pass-through” function that does not change the node or the state.

Parameters:

state – the (supposedly immutable) state object passed during the initial call.
node – the current node
ctx – the (supposedly immutable) dictionary with the global context passed during the initial call. In addition to normal dictionary methods, its values can be alternatively accessed as attributes (e.g. either ctx['value'] or ctx.value).
prepend – a function prepend(lst) which, when called, prepends the list of ast.AST objects to whatever is returned by the handler of the closest statement block that includes the current node. These nodes are not traversed automatically.
visit_after – a function of no arguments, which, when called, schedules to call the handler again on this node when all of its fields are traversed (providing that after calling it, the handler returns an ast.AST object and not a list or None). During the second call this parameter is set to None.
visiting_after – set to False during the normal (pre-order) visit, and to True during the visit caused by visit_after().
skip_fields – a function of no arguments, which, when called, orders the walker not to traverse this node’s fields.
walk_field – a function walk_field(state, value, block_context=False) -> (new_state, new_value), which traverses the given field value. If the value contains a list of statements, block_context must be set to True, so that prepend could work correctly.

Returns:

must return a tuple (new_state, new_node), where new_node is one of:

None, in which case the corresponding node will be removed from the parent list or the parent node field.
The passed node (unchanged). By default, its fields will be traversed (unless skip_fields() is called).
A new ast.AST object, which will replace the passed node in the AST. By default, its fields will not be traversed, and the handler must do it manually if needed (by calling walk_field()).
If the current node is an element of a list, a list of ast.AST objects can be returned, which will be spliced in place of the node. Same as in the previous case, these new nodes will not be automatically traversed.

peval.ast_inspector(handler)#

A shortcut for ast_walker() which does not transform the tree, but only collects data. Therefore:

the resulting walker returns only the resulting state;
the handler must return only the new (or the unchanged given) state instead of a tuple (new_state, new_node);
walk_field has the signature walk_field(state, value, block_context=False) -> new_state.

peval.ast_transformer(handler)#

A shortcut for ast_walker() with no changing state. Therefore:

the resulting walker has the signature def walker(node, ctx=None) and returns the transformed AST tree;
the handler must return only the transformed node instead of a tuple (new_state, new_node);
walk_field has the signature walk_field(value, block_context=False) -> new_value.

peval.try_peval_expression(node, bindings)#: Try to partially evaluate the AST expression node using the dictionary bindings. Returns a pair (evaluated, result), where evaluated is a boolean and result is the evaulation result if evaluated is True, and an AST expression otherwise.

class peval.Dispatcher(handler_obj: Callable[[immutabledict, AST, immutabledict], Tuple[immutabledict, AST]], default_handler: Optional[Callable[[immutabledict, AST, immutabledict], Tuple[immutabledict, AST]]] = None)#

A dispatcher that maps a call to a group of functions based on the type of the first argument (hardcoded to be an AST node at the moment).

handler_obj can be either a function with the signature:

def handler(*args, **kwds)

or a class with the static methods:

@staticmethod
def handle_<tp>(*args, **kwds)

where <tp> is the name of the type that this function will handle (e.g., handle_FunctionDef for ast.FunctionDef). The class can also define the default handler:

@staticmethod
def handle(*args, **kwds)

If it is not defined, the default_handler value will be used (which must be a function with the same signature as above). If neither handle exists or default_handler is provided, a ValueError is thrown.

class peval.Function(tree: Union[AsyncFunctionDef, FunctionDef], globals_, closure_vals, compiler_flags: int)#

A wrapper for functions providing transformations to and from AST and simplifying operations with associated global and closure variables.

tree#: An AST tree of the function (including the header and decorators, if any)

globals#: A dictionary of global variables associated with the function.

closure_vals#: A dictionary of closure variables associated with the function.

bind_partial(*args, **kwds) → Function#: Binds the provided positional and keyword arguments and returns a new Function object with an updated signature.

eval() → Callable#: Evaluates and returns a callable function.

classmethod from_object(func: Callable, ignore_decorators: bool = False) → Function#: Creates a Function object from an evaluated function.

get_external_variables() → Dict[str, Any]#: Returns a unified dictionary of external variables for this function (both globals and closure variables).

get_source() → str#: Generates the function’s source code based on its tree.

replace(tree: Optional[FunctionDef] = None, globals_: Optional[Dict[str, Any]] = None) → Function#: Replaces the AST and/or globals and returns a new Function object. If some closure variables are not used by a new tree, adjusts the closure cells accordingly.