CODEWITHSUNDEEP
pythonmypytyping
surj
surj

Reputation: 4913

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

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:

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

Upvotes: 0

Views: 695

Answers (1)

Alex Waygood
Alex Waygood

Reputation: 7569

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

Upvotes: 1

Related Questions