New option in decorator to prepend positional arguments

src/decorator.py

@@ -72,6 +72,11 @@ def iscoroutinefunction(f):
     def isgeneratorfunction():
         return False
 
+try:  # python 3.3+
+    from inspect import signature
+except ImportError:
+    from funcsigs import signature
+
 
 DEF = re.compile(r'\s*def\s*([_\w][_\w\d]*)\s*\(')
 
@@ -108,21 +113,7 @@ def __init__(self, func=None, name=None, signature=None,
                     setattr(self, a, getattr(argspec, a))
                 for i, arg in enumerate(self.args):
                     setattr(self, 'arg%d' % i, arg)
-                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)
+                self.refresh_signature()
                 self.dict = func.__dict__.copy()
         # func=None happens when decorating a caller
         if name:
@@ -159,6 +150,25 @@ def update(self, func, **kw):
         func.__module__ = getattr(self, 'module', callermodule)
         func.__dict__.update(kw)
 
+    def refresh_signature(self):
+        """Update self.signature and self.shortsignature based on self.args,
+        self.varargs, self.varkw"""
+        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)
+
     def make(self, src_templ, evaldict=None, addsource=False, **attrs):
         "Make a new function from a given template and update the signature"
         src = src_templ % vars(self)  # expand name and signature
@@ -196,12 +206,15 @@ def make(self, src_templ, evaldict=None, addsource=False, **attrs):
 
     @classmethod
     def create(cls, obj, body, evaldict, defaults=None,
-               doc=None, module=None, addsource=True, **attrs):
+               doc=None, module=None, addsource=True, add_args=(), **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 add_args is not empty, these arguments will be prepended to the
+        positional arguments.
         """
         if isinstance(obj, str):  # "name(signature)"
             name, rest = obj.strip().split('(', 1)
@@ -219,15 +232,136 @@ def create(cls, obj, body, evaldict, defaults=None,
                 'return', 'return await')
         else:
             body = 'def %(name)s(%(signature)s):\n' + ibody
-        return self.make(body, evaldict, addsource, **attrs)
 
+        # Handle possible signature changes
+        sig_modded = False
+        if len(add_args) > 0:
+            # prepend them as positional args - hence the reversed()
+            for arg in reversed(add_args):
+                if arg not in self.args:
+                    self.args = [arg] + self.args
+                    sig_modded = True
+                else:
+                    # the argument already exists in the wrapped
+                    # function, nothing to do.
+                    pass
+        if sig_modded:
+            self.refresh_signature()
+
+        # make the function
+        func = self.make(body, evaldict, addsource, **attrs)
+
+        if sig_modded:
+            # delete this annotation otherwise inspect.signature
+            # will wrongly return the signature of func.__wrapped__
+            # instead of the signature of func
+            del func.__wrapped__
+
+        return func
+
+
+def _extract_additional_args(f_sig, add_args_names, args, kwargs):
+    """
+    Processes the arguments received by our caller so that at the end, args
+    and kwargs only contain what is needed by f (according to f_sig). All
+    additional arguments are returned separately, in order described by
+    `add_args_names`. If some names in `add_args_names` are present in `f_sig`,
+    then the arguments will appear both in the additional arguments and in
+    *args, **kwargs.
+
+    In the end, only *args can possibly be modified by the procedure (by removing
+    from it all additional arguments that were not in f_sig and were prepended).
+
+    So the result is a tuple (add_args, args)
+
+    :return: a tuple (add_args, args) where `add_args` are the values of
+        arguments named in `add_args_names` in the same order ; and `args` is
+        the positional arguments to send to the wrapped function together with
+        kwargs (args now only contains the positional args that are required by
+        f, without the extra ones)
+    """
+    # -- first the 'truly' additional ones (the ones not in the signature)
+    add_args = [None] * len(add_args_names)
+    for i, arg_name in enumerate(add_args_names):
+        if arg_name not in f_sig.parameters:
+            # remove this argument from the args and put it in the right place
+            add_args[i] = args[0]
+            args = args[1:]
+
+    # -- then the ones that already exist in the signature. Thanks,inspect pkg!
+    bound = f_sig.bind(*args, **kwargs)
+    for i, arg_name in enumerate(add_args_names):
+        if arg_name in f_sig.parameters:
+            add_args[i] = bound.arguments[arg_name]
 
-def decorate(func, caller, extras=()):
+    return add_args, args
+
+
+def _wrap_caller_for_additional_args(func, caller, additional_args):
+    """
+    This internal function wraps the caller so as to handle all cases
+    (if some additional args are already present in the signature or not)
+    so as to ensure a consistent caller signature.
+
+    :return: a new caller wrapping the caller, to be used in `decorate`
+    """
+    f_sig = signature(func)
+
+    # We will create a caller above the original caller in order to check
+    # if additional_args are already present in the signature or not, and
+    # act accordingly
+    original_caller = caller
+
+    # -- then create the appropriate function signature according to
+    # wrapped function signature assume that original_caller has all
+    # additional args as first positional arguments, in order
+    if not isgeneratorfunction(original_caller):
+        def caller(f, *args, **kwargs):
+            # Retrieve the values for additional args.
+            add_args, args = _extract_additional_args(f_sig, additional_args,
+                                                      args, kwargs)
+
+            # Call the original caller
+            return original_caller(f, *itertools.chain(add_args, args),
+                                   **kwargs)
+    else:
+        def caller(f, *args, **kwargs):
+            # Retrieve the value for additional args.
+            add_args, args = _extract_additional_args(f_sig, additional_args,
+                                                      args, kwargs)
+
+            # Call the original caller
+            for res in original_caller(f, *itertools.chain(add_args, args),
+                                       **kwargs):
+                yield res
+
+    return caller
+
+
+def decorate(func, caller, extras=(), additional_args=()):
     """
     decorate(func, caller) decorates a function using a caller.
     If the caller is a generator function, the resulting function
     will be a generator function.
+
+    You can provide additional arguments with `additional_args`. In that case
+    the caller's signature should be
+
+        `caller(f, <additional_args_in_order>, *args, **kwargs)`.
+
+    `*args, **kwargs` will always contain the arguments required by the inner
+    function `f`. If `additional_args` contains argument names that are already
+    present in `func`, they will be present both in <additional_args_in_order>
+    AND in `*args, **kwargs` so that it remains easy for the `caller` both to
+    get the additional arguments' values directly, and to call `f` with the
+    right arguments.
     """
+    if len(additional_args) > 0:
+        # wrap the caller so as to handle all cases
+        # (if some additional args are already present in the signature or not)
+        # so as to ensure a consistent caller signature
+        caller = _wrap_caller_for_additional_args(func, caller, additional_args)
+
     evaldict = dict(_call_=caller, _func_=func)
     es = ''
     for i, extra in enumerate(extras):
@@ -246,11 +380,12 @@ def decorate(func, caller, extras=()):
     if generatorcaller:
         fun = FunctionMaker.create(
             func, "for res in _call_(_func_, %s%%(shortsignature)s):\n"
-                  "    yield res" % es, evaldict, __wrapped__=func)
+                  "    yield res" % es, evaldict,
+            add_args=additional_args, __wrapped__=func)
     else:
         fun = FunctionMaker.create(
             func, "return _call_(_func_, %s%%(shortsignature)s)" % es,
-            evaldict, __wrapped__=func)
+            evaldict, add_args=additional_args, __wrapped__=func)
     if hasattr(func, '__qualname__'):
         fun.__qualname__ = func.__qualname__
     return fun