pythonpython-typingpython-decoratorspyright

Type Annotations and Docstrings for Decorator


I'm trying to add type annotations for my python decorator, but its proving to be really difficult to get right.

Here is my decorator:

RT = TypeVar('RT')
P = ParamSpec('P')

def meraki_dashboard_setup(function: Callable[P, RT]):
    """
    Adds both a "meraki_key" and a "dashboard" argument to a decorated function. If the
    "meraki_key" argument is provided, the function will setup the dashboard object using the
    key. If the "dashboard" argument is provided, the function will use it directly.
    """
    @wraps(function)
    def wrapper(
        dashboard: Optional[DashboardAPI] = None,
        meraki_key: Optional[MerakiKey] = None,
        *args: P.args, **kwargs: P.kwargs
    ):
        if not dashboard and meraki_key:
            dashboard = setup_meraki_dashboard(meraki_key)
        if dashboard and meraki_key:
            raise ValueError(
                "Both a dashboard object and a Meraki key were provided. Only one should be used."
            )
        if not dashboard and not meraki_key:
            raise ValueError(
                "Neither a dashboard object nor a Meraki key were provided."
            )
        if not dashboard:
            raise ValueError("Unable to setup the Meraki dashboard object")

        kwargs["dashboard"] = dashboard
        return function(*args, **kwargs)

    return wrapper

And here is my decorated function

@meraki_dashboard_setup
def get_organization_networks(
    organization_id: str,
    **kwargs: Any
):
    """
    Fetch all networks from a specific organization.
    """
    dashboard = cast(DashboardAPI, kwargs.get("dashboard"))
    return cast(
        List[MerakiOrgNetwork],
        dashboard.organizations.getOrganizationNetworks(organization_id)
    )

My problem is that it "almost" works. When I use the get_organization_networks function, the function will make sure that the parameters "dashboard", "meraki_key" and "organization_id" have the correct type, and it shows those parameters when you hover over the function.

However, since the decorator is inferring the output type of the decorator as the inferred output type of the "wrapper" function after being decorated with "wrapped", the only hint I get through the linter is this

(function) get_organization_networks: _Wrapped[(organization_id: str, **kwargs: Any), List[MerakiOrgNetwork], (dashboard: DashboardAPI | None = None, meraki_key: MerakiKey | None = None, organization_id: str, **kwargs: Any), List[MerakiOrgNetwork]]

This hint basically consists of _Wrapped[P, RT, Mix of wrapper arguments + P, RT]. This is almost perfect, but I would love to get just this (like in a regular function):

(function) get_organization_networks(dashboard: DashboardAPI | None = None, meraki_key: MerakiKey | None = None, organization_id: str, **kwargs: Any) -> List[MerakiOrgNetwork]
Fetch all networks from a specific organization.

Seems like functools's wraps decorator is not able to finish copying the metadata for the wrapper into the function, causing no docstring from the decorated function to be copied, as well as making the function signature appear all wonky. One way I found of forcing the wrapper to by typed properly is by adding an output type to the meraki_dashboard_setup function

def meraki_dashboard_setup(function: Callable[P, RT]) -> Callable[P, RT]:
   ...

This returns the following hint:

(function) def get_organizations(**kwargs: Any) -> List[MerakiOrg]
Fetch all organizations from the Meraki Dashboard API

I've investigated how to get this working, but I cannot figure it out for the life of me. Maybe something to do with ParamSpec? I could technically leave it with my initial approach and it will work, but I really would like to have docstrings and a proper function signature. Its going to drive me insane if it doesn't. What could I do?


Solution

  • You want Concatenate:

    (playground)

    def meraki_dashboard_setup(function: Callable[P, RT]) \
      -> Callable[Concatenate[DashboardAPI | None, MerakiKey | None, P], RT]:
      ...
    

    The decorated function will then have the following signature hint:

    (function) def get_organization_networks(DashboardAPI | None, MerakiKey | None, organization_id: str, **kwargs: Any) -> None

    For more precise signatures, you would have to use a Protocol:

    (playground)

    class SetupFunction(Protocol[P, RT]):
        def __call__(
            self,
            dashboard: DashboardAPI | None = None, meraki_key: MerakiKey | None = None,
            *args: P.args, **kwargs: P.kwargs
        ) -> RT:
            ...
    
    def meraki_dashboard_setup(function: Callable[P, RT]) -> SetupFunction[P, RT]:
        ...
    

    However, that comes with a trade-off: The signature hint wouldn't be as nice.

    SetupFunction[(organization_id: str, **kwargs: Any), None]

    To add on top of that, putting a new parameter in the middle of P.args and P.kwargs (e.g. (*args: P.args, foo: str, **kwargs: P.kwargs)) is invalid.

    Regardless, Ctrl + Shift + Space (Trigger Parameter Hints) will always show the expected results:

    (dashboard: DashboardAPI | None = None, meraki_key: MerakiKey | None = None, organization_id: str, **kwargs: Any) -> None

    I recommend mixing @overload with Protocol:

    (playground)

    class SetupFunction(Protocol[P, RT]):
        @overload
        def __call__(self, dashboard: DashboardAPI, *args: P.args, **kwargs: P.kwargs) -> RT: ...
        @overload
        def __call__(self, meraki_key: MerakiKey, *args: P.args, **kwargs: P.kwargs) -> RT: ...
    

    (dashboard: DashboardAPI, organization_id: str, **kwargs: Any) -> None; (meraki_key: MerakiKey, organization_id: str, **kwargs: Any) -> None