Return type depending on literal parameter

BeeOnRope 0
Viewed 34

When typing my python functions PEP-484 style, how can I make the return type depend on the literal value of a parameter?

E.g., take this trivial function which adds 1 to the first element in the list, and optionally returns the list if the second parameter is True.

def add_one(x: list[int], return_result: bool):
    x[0] += 1
    if result_result:
        return x

By default the return type will be inferred to be None | list[int] but I want it to be just None if literal False is passed or list[int] if literal True is passed as the second parameter.

python-typing python-3.x python
1 Answers

As @joel stated briefly in his comment, Literal and overload in fact do the trick together:

from typing import overload, Literal, Optional


@overload
def add_one(x: list[int], return_result: Literal[False]) -> None:
    ...


@overload
def add_one(x: list[int], return_result: Literal[True]) -> list[int]:
    ...


def add_one(x: list[int], return_result: bool) -> Optional[list[int]]:
    x[0] += 1
    if return_result:
        return x
    return None  # required by `mypy`


a = add_one([0], True)
a.append(2)  # no issues

b = add_one([0], False)
b.append(2)  # `mypy` error: "None" has no attribute "append"
Daniil Fajnberg0
Related