Type checking functions where return type is always None if an argument is None

[JukkaL]

Some time ago I encountered another issue in production code that we may need to solve before implementing strict None type checking (#357). Consider a function like this:

def convert(s: Optional[str]) -> Optional[int]:
    if s is None:
        return None
    return int(s)

Now code like this would be rejected according to PEP 484 rules, as the return type unconditionally includes None:

convert('2') + 1  # Error: can't add None and int

We could use overloading or single dispatch to model this, but this would involve pretty major refactoring for existing code and probably some performance overhead.

Alternatively, we could model this as a generic function, but currently there's no syntax to model the above case. Here is a strawman proposal:

from typing import OptionalTypeVar

Maybe = OptionalTypeVar('Maybe')

def convert(s: Maybe[str]) -> Maybe[int]:
    if s is None:
        return None
    return int(s)

OptionalTypeVar would work a bit like TypeVar with values, such as AnyStr, with a slight twist: the first value would always be None, and the second value is the index value, and this could have multiple different values in different positions. So Maybe[T] would be replaced with None or T, but all instances of Maybe would get either the first or second value everywhere in lock step fashion. The above function would thus be equivalent to this from type checking perspective (assuming we had overloading):

@overload
def convert(s: None) -> None:
    if s is None:
        return None
    return int(s)

@overload
def convert(s: str) -> int:
    if s is None:
        return None
    return int(s)

We could generalize this to arbitrary collections of types. I can't come up with a good syntax, but here is the first thing that comes to mind:

from typing import Alternate

def convert(s: Alternate[None, str]) -> Alternate[None, int]:
    ...  # Same as above

The semantics would be the same as above. Now we could define a function that maps ints to string and vice versa:

from typing import Alternate

def switch(s: Alternate[int, str]) -> Alternate[str, int]:
    ... 

Here the issue is that we can't have multiple alternate variables that could vary independently. Maybe there would be a way to define additional Alternate variants with different names, similar to TypeVar.

The latter could also deal with True and False polymorphism assuming they would be considered subtypes of bool. Here is an example that is similar to some examples in the std library (in subprocess, I think):

class Stream(Generic[Alternate[str, bytes]]):
    def __init__(self, unicode: Alternate[True, False]) -> None: ...

    def stream(self) -> IO[Alternate[str, bytes]]: ...

Interestingly, we could replace AnyStr with this definition instead of using TypeVar:

AnyStr = Alternate[str, bytes]

[o11c]

I've thought a bunch about how C++ does it.

I do think the right thing is to model it internally as a generic, though I'm not sure it that should be explicit, or implied because Union.

Suggested syntax:

def convert(s: Optional[str]) -> type(int() if s else None):
    if s is None:
        return None
    return int(s)

[o11c]

For True, False, etc. I believe the best syntax is Constant[True] as a type, and Constant(True) if you ever want it as a value (to avoid type annotations)

[jhance]

Its worth keeping in mind that f(x: Maybe[int], y: Maybe[int]) -> Maybe[int]' would probably want to return None if either x or y are None. I don't think this is quite the same as what you are proposing.

Such functions are also silly in a design where you have Optional (because you can just lift them "into Optional"), so the only reason to support such a function is honestly for legacy reasons.

[markshannon]

Why is

def convert(s: Optional[str]) -> Optional[int]:
    if s is None:
        return None
    return int(s)

illegal according to PEP 484? It looks perfectly legal to me.
def convert(s: Optional[str]) -> Optional[int] is a valid type and the function body is obviously type correct.

[rwbarton]

convert itself is fine, but the type def convert(s: Optional[str]) -> Optional[int] is not specific enough to allow the call convert('2') + 1 to type check.

[markshannon]

What level of inference is it reasonable to expect from a checker?
It would be nice to just write

def convert(s: Optional[str]):
    if s is None:
        return None
    return int(s)

and have the function type correctly inferred from the source.
For stubs, @overload already exists.

[gvanrossum]

In general, inference across functions is not something I'd like to tackle yet. This would imply whole-program analysis which is the graveyard of over-ambitious Python type checking projects.

However in this particular case you could argue that a sufficiently smart type checker itself could take the definition of convert() and infer the overloaded rewrite without help. That's nice but I don't want PEP 484 to require it.

So for now you'll have to write it using an overload. Hopefully allowing overloads in non-stub files will make this less painful. (#1136, also python/typing#175)

[JelleZijlstra]

Jukka's original example can be expressed with overloads, even outside stub files. Some of the ideas in this issue are broader, but since the original use case has been addressed, I don't think it's worth it to keep this issue open.

[MetRonnie]

The correct way of using overloads for the Jukka's example would be

from typing import Optional, overload


@overload
def convert(s: str) -> int: ...


@overload
def convert(s: None) -> None: ...


def convert(s: Optional[str]) -> Optional[int]:
    if s is None:
        return None
    return int(s)

Edit: see also #6113 (comment) for when using Literal[True/False] with bool