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.
Helper functions#
- peval.getsource(func: Callable) str #
Returns the source of a function
func
. Falls back toinspect.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 usingfunctools.lru_cache
with the providedmaxsize
(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 aDispatcher
instance; seeDispatcher
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)
, wherestate
is the same object which was passed as the corresponding parameter. Does not mutatenode
.
handler
will be invoked for every node during the AST traversal (depth-first, pre-order). Thehandler
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']
orctx.value
).prepend – a function
prepend(lst)
which, when called, prepends the list ofast.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 orNone
). During the second call this parameter is set toNone
.visiting_after – set to
False
during the normal (pre-order) visit, and toTrue
during the visit caused byvisit_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 toTrue
, so thatprepend
could work correctly.
- Returns:
must return a tuple
(new_state, new_node)
, wherenew_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 (unlessskip_fields()
is called).A new
ast.AST
object, which will replace the passednode
in the AST. By default, its fields will not be traversed, and the handler must do it manually if needed (by callingwalk_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 signaturewalk_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 signaturewalk_field(value, block_context=False) -> new_value
.
- peval.try_peval_expression(node, bindings)#
Try to partially evaluate the AST expression
node
using the dictionarybindings
. Returns a pair(evaluated, result)
, whereevaluated
is a boolean andresult
is the evaulation result ifevaluated
isTrue
, 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
forast.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 neitherhandle
exists ordefault_handler
is provided, aValueError
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.