Python typing for function that returns None if list arg only contains None

Question

I'm working with a function that is a bit like this (super simplified, as an example):

def foo(*stuff: None | int):
    stuff_not_none = [x for x in stuff if x is not None]
    if len(stuff_not_none) is 0:
        return None
    return sum(stuff_not_none)

If I call the function using:

• foo(*[1, 2, 3]), I'd want the return type to be inferred to int.
• foo(*[None, None]), I'd want the return type to be inferred to None.
• foo(*[1, None]), the dream would be inferred to int, but ok if None | int.

I've tried with generics / overloads, but I couldn't figure out this puzzle. How can I achieve this?

Answer 1

The solution:

from typing import overload

@overload
def foo(*stuff: None) -> None: ...  # type: ignore[misc]

@overload
def foo(*stuff: int | None) -> int: ...

def foo(*stuff: int | None) -> int | None:
    stuff_not_none = [x for x in stuff if x is not None]
    if len(stuff_not_none) is 0:
        return None
    return sum(stuff_not_none)
    
reveal_type(foo(None, None))  # revealed type is None
reveal_type(foo(1, 2, 3))  # revealed type is int
reveal_type(foo(None, 2, None, 4))  # revealed type is int
foo('a', 'b')  # error: no matching overload

Mypy hates this kind of thing, because the overloads overlap. But you'll find that if you add a type: ignore comment in the right place, it's perfectly able to infer the correct types anyway. (I'm a typeshed maintainer, and we do this kind of thing at typeshed all the time.)

Note that the order of the overloads is very important: type checkers will always try the first overload first, and then, only if that doesn't match, will they try the second overload. This is how we get the int revealed type when we pass in a mixture of ints and Nones: the first overload doesn't match, because of the presence of ints, so the type checker is forced to try the second overload.

Mypy playground demo: https://mypy-play.net/?mypy=latest&python=3.10&gist=ff07808e0a314208fdfa6291dcf9f717