Search code examples
pythonpython-typing

Python typing annotation return value annotation based if function argument being a list or not

If I have

   def get(
        ids: str | list[str] | int | list[int],
    ) -> float | list[float]:

Is there a way to specify in return value annotation that a list of floats is output only when the input ids is a list of strs or ints?

Solution

  • One way to do this is by using the @overload decorator to create 2 extra function signatures, one for int | str that returns float and one for list[str] | list[int] that returns list[float] and then have the actual function definition like you have now.

    from typing import overload


@overload
def test(
    ids:str | int,
) -> float:...

@overload
def test(
    ids:list[str] | list[int],
) -> list[float]:...

def test(
    ids: str | list[str] | int | list[int],
) -> float | list[float]:...