Get proper intellisense with overloaded multimethods in python

Question

The following code works as expected:

from multimethod import multimethod

@multimethod
def foo(a: int):
    print(f"only {a}")

@multimethod
def foo(a: int, b: int):
    print(f"{a} and {b}")

foo(1)
foo(1, 2)

Output:

only 1
1 and 2

However, the Intellisense in VSCode for the function arguments look like this

I know that this is caused by the multimethod decorator expecting any number of arguments for it to decide which overload of the function to call, however what I want is something like this:

with multiple argument lists.

Does anyone know how to achieve this?

Answer 1

With the help of @juanpa.arrivillaga I solved this by creating a stub file alongside the file that contains the functions, this worked perfectly. Although my actual code is different from this, the stub file for the example foo function would look like this:

foo.pyi

from typing import overload

@overload
def foo(a: int): ...
@overload
def foo(a: int, b: int): ...

This will provide whatever editor you're using (i.e. VSCode) with the information to display parameter info of both overloads.