Help With Decorator Parameter Hints

[Logon27]

Hello. I am manually using a python decorator to wrap a function. Inside the decorator I have proper functools.wraps calls for updating the metadata. However, Pylance gives me this terrible "_Wrapped" syntax...

When it should be shown like...

Is there any way to get Pylance to display the correct function parameter hints? I have tried a bunch of different methods and nothing seems to work.

[erictraut]

Could you provide a minimal, self-contained code sample (in text form) that shows the problem you're trying to solve? It's difficult to answer your question from partial screen shots.

[erictraut]

I'm not sure I fully understand what this code is supposed to do. Is it possible that you omitted a return statement in the first Dense function?

I've taken a stab at annotating the code, but I needed to make some assumptions. Let me know if my assumptions were incorrect.

from typing import Any, Callable, ParamSpec
import functools

P1 = ParamSpec("P1")

def debug_decorator(
    dense_debug: Callable[P1, tuple[Callable[..., Any], Callable[..., Any]]]
):
    @functools.wraps(dense_debug)
    def dense(*args: P1.args, **kwargs: P1.kwargs):
        init_fun_debug, apply_fun_debug = dense_debug(*args, **kwargs)

        @functools.wraps(init_fun_debug)
        def init_fun(input_shape: tuple[int, ...]):
            return input_shape

        @functools.wraps(apply_fun_debug)
        def apply_fun(inputs: tuple[int, ...]):
            return inputs

        return init_fun, apply_fun

    return dense

def dense(out_dim: int):
    def init_fun(input_shape: tuple[int, ...]):
        return input_shape

    def apply_fun(inputs: tuple[int, ...]):
        return inputs

    return init_fun, apply_fun

dense2 = debug_decorator(dense)

[Logon27]

That example was me cutting all out dozens of lines of code to give a dummy example. Let's just simplify to a trivial example so we are on the same page here. An example function that just takes in a number and returns it. An example_decorator that just wraps the function and prints the value.

import functools

def example_decorator(example_debug):

    @functools.wraps(example_debug)
    def new_example(*args, **kwargs):
        print(f"my number: {args[0]}")
        example_debug(*args, **kwargs)

    return new_example

def example(my_number: int):
    return my_number

example2 = example_decorator(example)

Edit: I changed the inner decorator function to something unique just so we are confusing the function scope here.
Edit2: Whoops. Removed the decorator annotation since we are calling the decorator directly.

[erictraut]

You can use a ParamSpec to preserve the signature of the decorated function:

P = ParamSpec("P")

def example_decorator(example_debug: Callable[P, Any]):
    @functools.wraps(example_debug)
    def new_example(*args: P.args, **kwargs: P.kwargs):
        print(f"my number: {args[0]}")
        example_debug(*args, **kwargs)

    return new_example

[erictraut]

If you want to preserve the return type as well, you'll need to return it from within your new_example inner function.

P = ParamSpec("P")
R = TypeVar("R")

def example_decorator(example_debug: Callable[P, R]) -> Callable[P, R]:
    @functools.wraps(example_debug)
    def new_example(*args: P.args, **kwargs: P.kwargs) -> R:
        print(f"my number: {args[0]}")
        return example_debug(*args, **kwargs)

    return new_example

[erictraut]

For more details about ParamSpec, refer to PEP 612.

[Logon27]

The new parameter hint...

Like yes, the signature has improved slightly. But it is still a far cry from the original parameter hint...

For any non-trivial example, this "_Wrapped" syntax is still pretty hard to read and does not respect new line formatting and color differentiation. Is there anything I can do to improve this even more?

[Logon27]

P = ParamSpec("P")
R = TypeVar("R")

def example_decorator(example_debug: Callable[P, R]) -> Callable[P, R]:
    @functools.wraps(example_debug)
    def new_example(*args: P.args, **kwargs: P.kwargs) -> R:
        print(f"my number: {args[0]}")
        return example_debug(*args, **kwargs)

    return new_example

@erictraut Bizarre, only your second example with both the ParamSpec and TypeVar produces the correct parameter hint. Your first example doesn't give the same level of parameter hints. Your second example (which works)...

versus your first example snippet...

[Logon27]

I have to get off for the night. Thank you Eric, I appreciate all the help. I know nothing about these parameter specs so I will have to read up on it. A tad annoying that I have to specifically type hint my code with this format in order to get the proper parameter hints. But it is better than having no working solution at all.