Possibility to add additional arguments to the decorated function

[smarie]

Hi there,

I'm a big fan of your library and I used it in many of mine :).

I recently came in pytest-steps with the need to

• add arguments to a decorated function
• but tolerate that these arguments can already be there

The use-case is that the decorated function is a pytest test function, that can or can not already have the "request" argument in its signature, but I need the wrapped function to have it always. I came up with the following hack, maybe this is something you would be interested in ?

class MyFunctionMaker(FunctionMaker):
    """
    Overrides FunctionMaker so that additional arguments can be inserted in the resulting signature.
    """

    @classmethod
    def create(cls, obj, body, evaldict, defaults=None,
               doc=None, module=None, addsource=True, add_args=None, **attrs):
        """
        Create a function from the strings name, signature and body.
        evaldict is the evaluation dictionary. If addsource is true an
        attribute __source__ is added to the result. The attributes attrs
        are added, if any.
        """
        if isinstance(obj, str):  # "name(signature)"
            name, rest = obj.strip().split('(', 1)
            signature = rest[:-1]  # strip a right parens
            func = None
        else:  # a function
            name = None
            signature = None
            func = obj
        self = cls(func, name, signature, defaults, doc, module)
        ibody = '\n'.join('    ' + line for line in body.splitlines())
        caller = evaldict.get('_call_')  # when called from `decorate`
        if caller and iscoroutinefunction(caller):
            body = ('async def %(name)s(%(signature)s):\n' + ibody).replace(
                'return', 'return await')
        else:
            body = 'def %(name)s(%(signature)s):\n' + ibody

        # --- HACK part 1 -----
        if add_args is not None:
            for arg in add_args:
                if arg not in self.args:
                    self.args = [arg] + self.args
                else:
                    # the argument already exists in the wrapped function, no problem.
                    pass

            # update signatures (this is a copy of the init code)
            allargs = list(self.args)
            allshortargs = list(self.args)
            if self.varargs:
                allargs.append('*' + self.varargs)
                allshortargs.append('*' + self.varargs)
            elif self.kwonlyargs:
                allargs.append('*')  # single star syntax
            for a in self.kwonlyargs:
                allargs.append('%s=None' % a)
                allshortargs.append('%s=%s' % (a, a))
            if self.varkw:
                allargs.append('**' + self.varkw)
                allshortargs.append('**' + self.varkw)
            self.signature = ', '.join(allargs)
            self.shortsignature = ', '.join(allshortargs)
        # ---------------------------

        func = self.make(body, evaldict, addsource, **attrs)

        # ----- HACK part 2
        if add_args is not None:
            # delete this annotation otherwise the inspect.signature method relies on the wrapped object's signature
            del func.__wrapped__

        return func


def my_decorate(func, caller, extras=(), additional_args=None):
    """
    A clone of 'decorate' with the possibility to add additional args to the function signature.
    Additional arguments will be positional and will sit at the beginning of 
    """
    evaldict = dict(_call_=caller, _func_=func)
    es = ''
    for i, extra in enumerate(extras):
        ex = '_e%d_' % i
        evaldict[ex] = extra
        es += ex + ', '
    fun = MyFunctionMaker.create(
        func, "return _call_(_func_, %s%%(shortsignature)s)" % es,
        evaldict, add_args=reversed(additional_args or ()), __wrapped__=func)
    if hasattr(func, '__qualname__'):
        fun.__qualname__ = func.__qualname__
    return fun

Here is how I use it in a complex case. One additional argument ________step_name_ is guaranteed to not be in the function signature, while the 'request' argument may or may not be there

# -- first create the logic
def _execute_step_with_monitor(step_name, request, *args, **kwargs):
    # ...

# -- then create the appropriate function signature according to wrapped function signature
if 'request' not in f_sig.parameters:
    # easy: we can add it explicitly in our signature
    def step_function_wrapper(f, ________step_name_, request, *args, **kwargs):
        """Executes a step with the execution monitor for this pytest node"""
        _execute_step_with_monitor(________step_name_, request, *args, **kwargs)
else:
    # harder: we have to retrieve the value for request. Thanks, inspect package !
    def step_function_wrapper(f, ________step_name_, *args, **kwargs):
        """Executes a step with the execution monitor for this pytest node"""
        request = f_sig.bind(*args, **kwargs).arguments['request']
        _execute_step_with_monitor(________step_name_, request, *args, **kwargs)

# decorate it so that its signature is the same than test_func, with just an additional argument for test step
# and if needed an additional argument for request
wrapped_test_function = my_decorate(test_func, step_function_wrapper,
                                    additional_args=['________step_name_', 'request'])

If you think that this feature is interesting, I can make a Pull Request.

[smarie]

Following conversation in #58 here is hopefully a simpler explanation:

Basically my decorator wraps user-created test functions that are compliant with pytest (in the example below, test_func). Therefore these test functions can add the special argument 'requests' in their signature if they wish pytest to inject that object into them. But they are not obliged to, they can also live without it.

The issue is that in my function wrapper, I absolutely need that 'requests' object to be injected - indeed I use it in the wrapper that I create. So I have to add the argument to the signature if it is not already present, or leave it if it is. For this purpose I improved your decorate method to add an additional_args argument (for the record, decorate is the single method I always use when I use your package - I find it handy and simple to use. I then create the decorator itself 'by hand' leveraging it).

# The function to wrap is 'test_func'. We first extract its signature
f_sig = signature(test_func)

# -- first define the logic once and for all
def _inner_logic(f, request, *args, **kwargs):
    # ... (my business here)

# -- then define the appropriate function wrapper according to wrapped function signature
if 'request' not in f_sig.parameters:
    # easy: we can add 'request' explicitly in our signature
    def dependency_mgr_wrapper(f, request, *args, **kwargs):
        return _inner_logic(f, request, *args, **kwargs)
else:
    # harder: we have to retrieve the value for request from args. Thanks, inspect package !
    def dependency_mgr_wrapper(f, *args, **kwargs):
        request = f_sig.bind(*args, **kwargs).arguments['request']
        return _inner_logic(f, request, *args, **kwargs)

# Finally wrap the test function and add the 'request' argument
wrapped_test_function = decorate(test_func, dependency_mgr_wrapper, additional_args=['request'])

[smarie]

Note: I'm also quite sure that one day someone will come up with the need to remove an argument too :) but let's leave it for later

[smarie]

I found an even better way to use this. With the new PR code, the above example would be much simpler to write:

def dependency_mgr_wrapper(f, request, *args, **kwargs):
    # ... (my business here)

# Finally wrap the test function and add the 'request' argument
wrapped_test_function = decorate(test_func, dependency_mgr_wrapper, additional_args=['request'])

Basically I put the signature checking steps in the decorate method, so that the behaviour is consistent whatever the usage case: the signature of the caller provided by the user should always be

def caller(f, <all named additional args in order>, *args, **kwargs):

If some additional args are already part of the wrapped function signature they will arrive twice. That way it is easy for users to call the inner function with all its required arguments, by calling f(*args, **kwargs)

[smarie]

For the record - this dynamic function creation with possibility to change the signature is now in makefun.