python_typing_optional_literal_pos_or_kw_param.md

original post: scipy/scipy-stubs#153 (comment)

---

Annotating a function whose return type depends on an optional positional- or keyword-only parameter in Python

Overloads are roughly similar to pattern matching¹, but from the perspective of the caller.

Consider this function for example

def f(x=0, b=False):
    y = (2 << x) - 1
    return x if b else str(x)

For simplicity let's only assume that we only accept builtin types, i.e. x: int and b: bool. The outcome depends out the literal value of b, so the output type depends on whether b is Literal[False] or Literal[True].

Because both x and b can be passed either positionally or by keyword, we're dealing with four possible call signatures:

1. int: f() or f(x) or f(x, False) (positional)
2. int: f(b=False) or f(x, b=False) (keyword)
3. str: f(b=True) or f(x, b=True) (keyword)
4. str: f(x, True) (positional)

We can annotate each of these using overloads:

from typing import Literal as L, overload

# 1.
@overload
def f(x: int = 0, b: L[False] = False) -> int: ...

# 2. (the `*` indicates that `b` is passed as keyword argument)
@overload
def f(x: int = 0, *, b: L[False]) -> int: ...

# 3.
@overload
def f(x: int = 0, *, b: L[True]) -> str: ... 

# 4. (note that `x` is required in this case, so it has no default)
@overload
def f(x: int, b: L[True]) -> str: ... 

Note that the second overload is actually redundant: The first one will already match f(b=False) or f(x, b=False), and both overloads have the same return type. So that leaves us with

@overload
def f(x: int = 0, b: L[False] = False) -> int: ...
@overload
def f(x: int = 0, *, b: L[True]) -> str: ... 
@overload
def f(x: int, b: L[True]) -> str: ... 

In this case the order doesn't matter. That's because none of the overloads overlap. If you want, you could equivalently write this as e.g.

@overload
def f(x: int, b: L[True]) -> str: ... 
@overload
def f(x: int = 0, *, b: L[True]) -> str: ... 
@overload
def f(x: int = 0, b: L[False] = False) -> int: ...

The order is just a matter of personal preference in this case. But whathever ordering you choose, try to be consistent with it.

Footnotes

1. Overloads actually arent't the same as pattern matching. But for this purpose, we can pretend that they are. ↩