pythonpython-decoratorsfunctools

How can I change the parameter signature in the help docs of a decorated function?


I have a class that includes a method that takes two parameters, a and b, like so:

class Foo:
  def method(self, a, b):
    """Does something"""
    x, y, z = a.x, a.y, b.z
    return do(x, y, z)

When running help(Foo) in the Python REPL, you would see something like this:

class Foo(builtins.object)
 |  Methods defined here:
 |  
 |  method(self, a, b)
 |      Does something

I have updated the class with a decorator that manipulates the arguments passed into the method so instead of taking a and b and unpacking them, the decorator will do the unpacking and pass x, y, and z into the method:

def unpack(fn):
  def _unpack(self, a, b):
    x, y, z = a.x, a.y, b.z
    return fn(self, x, y, z)
  return _unpack

class Foo:
  @unpack
  def method(self, x, y, z):
    """Does something"""
    return do(x, y, z)

This works fine except that the help string isn't very helpful anymore:

class Foo(builtins.object)
 |  Methods defined here:
 |  
 |  method = _unpack(self, a, b)

The standard way to fix this is to use functools.wraps:

from functools import wraps

def unpack(fn):
  @wraps(fn)
  def _unpack(self, a, b):
    x, y, z = a.x, a.y, b.z
    return fn(self, x, y, z)
  return _unpack

And that works great, too, except that it shows the method as taking x, y, and z as arguments, when really it still takes a and b (that is, 2 arguments instead of 3), so the help doc is a bit misleading:

class Foo(builtins.object)
 |  Methods defined here:
 |  
 |  method(self, x, y, z)
 |      Does something

Is there a way to modify the function wrapper so it correctly grabs the doc string and other attributes as in functools.wraps but also shows the arguments accepted by the wrapper?

For example, I would like the help string to show:

class Foo(builtins.object)
 |  Methods defined here:
 |  
 |  method(self, a, b)
 |      Does something

even when the method is wrapped by unpack.

(I have examined the source code of functools.wraps but I could not figure out if that can copy the function signature or not.)


Solution

  • TL;DR

    Delete the __wrapped__ attribute that is added automatically by functools.wraps before returning the wrapper:

    def unpack(fn):
        @wraps(fn)
        def wrapper(self, a, b):
            x, y, z = a.x, a.y, b.z
            return fn(self, x, y, z)
        del wrapper.__wrapped__  # <-- force inspection to show `wrapper` signature
        return wrapper
    

    Details

    The built-in help function is just a a wrapper around pydoc.help, which in turn is just an instance of pydoc.Helper.

    Skipping over a few steps, it turns out that eventually for rendering the documentation for any kind of function, the TextDoc.docroutine method is called. And that in turn utilizes inspect.signature to retrieve the signature of the function in question (see here).

    All it does then is take the string representation of that Signature object, adds a bit more stuff, and glues the docstring underneath it.

    The docstring is essentially just the __doc__ attribute of the function object in question, which functools.wraps (or more precisely functools.update_wrapper) copies over from the "actual" function. This is why using the @wraps decorator solves the docstring issue.

    The problem for you is that update_wrapper also always adds the __wrapped__ attribute to the wrapper function and assigns the original function to that. Turns out that the inspect library knows this and is on the lookout for that attribute. The idea presumably being that you usually want to inspect the "actual" underlying function, not the wrapper, since you took the time to use update_wrapper in the first place.

    The inspect.unwrap method is used to drill down to the non-wrapped function before analyzing that and retrieving its parameter specification. Therefore the straightforward solution is to simply get rid of the __wrapped__ attribute on your wrapper function so that inspect will be forced to analyze that.

    It should be obvious, but the __wrapped__ attribute is added for a reason: To still have access to the original function after decoration. So while this does solve your problem, it may introduce other problems.


    PS

    Taking this one step further, you could theoretically keep the __wrapped__ attribute, but assign it a stub with any signature you want. This will "fool" the inspect tools into returning the signature of that stub, which will in turn end up in the auto-generated documentation from pydoc/help.

    For example:

    from collections.abc import Callable
    from functools import wraps
    from typing import Protocol, TypeVar
    
    R = TypeVar("R")
    T = TypeVar("T")
    
    
    class HasXY(Protocol):
        x: int
        y: str
    
    
    class HasZ(Protocol):
        z: float
    
    
    def unpack(fn: Callable[[T, int, str, float], R]) -> Callable[[T, HasXY, HasZ], R]:
        @wraps(fn)
        def wrapper(self: T, a: HasXY, b: HasZ) -> R:
            x, y, z = a.x, a.y, b.z
            return fn(self, x, y, z)
    
        def _signature(self, a: HasXY, b: HasZ):  # type: ignore[no-untyped-def]
            raise NotImplementedError
        __annotations__ = getattr(fn, "__annotations__", {})
        for name in ("self", "return"):
            if name in __annotations__:
                _signature.__annotations__[name] = __annotations__[name]
        wrapper.__wrapped__ = _signature  # type: ignore[attr-defined]
    
        return wrapper
    
    class Foo:
        @unpack
        def method(self, x: int, y: str, z: float) -> None:
            """Does something"""
            print(f"method({x=}, {y=}, {z=})")
    
    
    help(Foo)
    
    
    class Bar:
        x = 1
        y = "a"
    
    
    class Baz:
        z = 3.14
    
    
    foo = Foo()
    foo.method(Bar, Baz)
    

    Output:

    Help on class Foo in module __main__:
    
    class Foo(builtins.object)
     |  Methods defined here:
     |  
     |  method(self, a: __main__.HasXY, b: __main__.HasZ) -> None
     |      Does something
     |  
     |  ----------------------------------------------------------------------
     |  ...
    
    method(x=1, y='a', z=3.14)
    

    The code above will incidentally pass mypy --strict without errors.

    Again, this is somewhat hack-ish, in that that the method's __wrapped__ attribute is now just a reference to the _signature stub, so you will not be able to call the original method directly in any way.