How do I annotate a Python function to hint that it takes the same arguments as another function?

Question

Is there any Python type-hinting syntax to state that a function takes the same parameters (and parameter types) as another function? In particular this is useful for wrapping, e.g.,

async def do_stuff(
        param1: str,
        param2: int,
        param3: int = 14,
):
    ...

def run_async_thing(*args, **kwargs):  # <--- What can I put here to say 'takes args like `do_stuff`'?
    return asyncio.get_event_loop().run_until_complete(do_stuff(*args, **kwargs))

In this case, I would like to add type hinting to the run_async_thing function to identify that it expects the same argument types as the do_stuff function.

Is this possible, and if so, how?

The primary reason for wanting this is so that my tools (in particular PyCharm/IntellliJ IDEA) can figure out what arguments run_async_thing should expect/accept. If it helps with documentation that's a bonus, but this is mainly for tooling.

Answer 1

Define the parameters explicitly. You are unnecessarily generalizing the signature for run_async_thing:

def run_async_thing(p1: str, p2: int, p3: int):
    return asyncio.get_event_loop().run_until_complete(do_stuff(p1, p2, p3))

More generally, you can have run_async_thing take a single tuple (or other object) as an argument. For example:

async def do_stuff(t: Tuple[str, int, int]):
    ...

def run_async_thing(args: Tuple[str, int, int]):
   return asyncio.get_event_loop().run_until_complete(do_stuff(args))

The tuple type can be factored out:

StuffDoerArgs = Tuple[str, int, int]

async def do_stuff(t: StuffDoerArgs):
    ...

def run_async_thing(args: StuffDoerArgs):
    ...

Answer 2

You need to define the parameter types and return types. You can then copy the __annotations__ from one function to the other.

Example:

def abc(num: int, name: str) -> list:
    return [num, name]

print(abc.__annotations__)

Output: {'num': <class 'int'>, 'name': <class 'str'>, 'return': <class 'list'>}

Now we create another function:

def xyz(num, name):
    return [num, name]

print(xyz.__annotations__)

Output: {}

You can just copy over the __annotations__ output from one to the other.

xyz.__annotations__ = abc.__annotations__

So now:

print(xyz.__annotations__)

Output: {'num': <class 'int'>, 'name': <class 'str'>, 'return': <class 'list'>}

Answer 3

I ran into the same problem just now, and found out a way that's concise and works in Pylance/Pyright.

In some cases it's not plausible to replicate the other function's signature, as it may be very complex.

The solution to this is to type hint one function as another. Here's an example with the original askers code:

async def do_stuff(
        param1: str,
        param2: int,
        param3: int = 14,
):
    ...

def run_async_thing(*args, **kwargs):
    return asyncio.get_event_loop().run_until_complete(do_stuff(*args, **kwargs))

run_async_thing: do_stuff = run_async_thing

This would be picked up by Pylance/Pyright

Answer 4

If you find yourself consistently using the set of parameters for different functions you could consider grouping them with a class-

Instead of

async def do_stuff(
        param1: str,
        param2: int,
        param3: int,
):
    ...

def run_async_thing(p1: str, p2: int, p3: int):
    ...

You could do

class AsyncArgs:
    def __init__(self, p1: str, p2: int, p3: int):
        self.p1 = p1
        self.p2 = p2
        self.p3 = p3

    def get_p1(self):
        return self.p1

    def get_p2(self):
        return self.p2

    def get_p3(self):
        return self.p3

async def do_stuff(args: AsyncArgs):
    ...

def run_async_thing(args: AsyncArgs):
    ...

Answer 5

I don't know if it will work in JetBrains or not, but as suggested in a separate related answer, you can use a decorator with type annotations to achieve this, and since it is decorator function, you could also handle the annotations there for documentation.

import asyncio
from typing import TypeVar, Callable, Any
from typing_extensions import ParamSpec

class Stuff:
    name = 3
    age = 5

T = TypeVar('T')
V = TypeVar('V')
P = ParamSpec('P')

def wraps(x: Callable[P, Any]):
    def decorator(f: Callable[..., V]) -> Callable[P, V]:
        f.__annotations__ = {**x.__annotations__, "return": f.__annotations__["return"]}
        return f
    return decorator

async def do_stuff(param1: str, param2: int, param3: int = 14) -> Stuff:
    pass

@wraps(do_stuff)
def run_async_thing(*args, **kwargs):
    return asyncio.get_event_loop().run_until_complete(do_stuff(*args, **kwargs))

result = run_async_thing(param1="hello", param2="world")

When I tried this with PyLance, the result is that it picks up that run_async_thing accepts the arguments of do_stuff, and it also picks up that the return value is a Stuff object.

The above uses the typing_extensions module, which may not be ideal. If you want to avoid that and don't care about the return type being accurate, you could just use the whole function as the generic.

def wraps(x: T):
    def decorator(f) -> T:
        f.__annotations__ = x.__annotations__
        return f
    return decorator

With this, the return value comes through as being a coroutine, although it is not.

Another thing I've noticed is that setting f.__annotations__ doesn't seem to do much from what I can tell - at least pydoc still shows the method parameters as defined directly on the function.

Another note (your question didn't as much ask about this) is that there doesn't seem to be a way to annotate a Callable as accepting both parameters from itself as well as from another method that it wraps. Using Union[Callable[P, T], Callable[R,T]] where P and R are the ParamSpec types of the two methods and T is the return type of the wrapper seems to get close, except that positional arguments, if any, are off, and keyword args either match one set or the other but not both at the same time.