Function Helpers

Module: polygraphy.func

extend(extend_func)[source]

A decorator that uses the function it decorates to extend the function provided as a parameter.

This is best illustrated with an example:

def x(a0, a1, a2):
    rv0 = [a0, a1, a2]
    rv1 = None
    return rv0, rv1

@extend(x)
def y(rv0, rv1):
    rv0.append(-1)

# We can now call `y` as if it were `x`, and we will receive
# the return values from `x` after any modifications by `y`
rv0, rv1 = y(1, 2, 3)
assert rv0 == [1, 2, 3, -1]
assert rv1 is None

In this case, extend is essentially syntactic sugar for:

def y(a0, a1, a2):
    rv0, rv1 = x(a0, a1, a2)

    # Body of `y` from previous section
    rv0.append(-1)

    return rv0, rv1

If y does not return anything, or returns None, then extend will ensure that the return value of x is forwarded to the caller. This means that y will provide exactly the same interface as x.

If y returns something other than None, then its return value will be provided to the caller, and the return value of x will be discarded.

In some cases, it may be necessary to access the parameters of x. If y takes all of x’s parameters prior to its usual parameters (which are the return values of x), then all arguments to x will be forwarded to y. Note that if x modifies its arguments in place, then y will see the modified arguments, not the original ones. For example:

def x(x_arg):
    return x_arg + 1

def y(x_arg, x_ret):
    # `y` can now see both the input argument as well as the return value of `x`
    assert x_ret == x_arg + 1

assert y(5) == 6

NOTE: This function will automatically unpack tuples returned by the function being extended. Thus, the following implementation of x would behave just like the one mentioned above:

def x(a0, a1, a2):
    ret = (rv0, rv1)
    return ret # Tuple will be unpacked, and `y` still sees 2 parameters

NOTE: The decorated function must not use variadic parameters like *args or **kwargs.

Parameters:: extend_func (Callable) – A callable to extend.