Should have a way to annotate an Optional carryover method.

[euresti]

Edited with a better example.

Let's suppose you have a function that can take an Optional and returns something based on that or None if it's None.

from typing import overload, Optional

def bar(foo: Optional[int]) -> Optional[str]:
   if foo is None:
       return None
   else:
       return str(foo)

x = None  # type: Optional[int]

reveal_type(bar(None))  # Line 18
reveal_type(bar(x))
reveal_type(bar('foo'))

There's one problem with this being annotated this way. If you pass a Non-Optional you get returned an Optional[str] even though it's not possible. So I thought the new @overload stuff might help here but it doesn't:

from typing import overload, Optional

@overload
def bar(foo: int) -> str:
    ...

@overload
def bar(foo: Optional[int]) -> Optional[str]:
    ...

def bar(foo: Optional[int]) -> Optional[str]:
   if foo is None:
       return None
   else:
       return str(foo)

x = None  # type: Optional[int]

reveal_type(bar(None))
reveal_type(bar(x))
reveal_type(bar('foo'))

Sadly that gives me an error and the wrong types:

tt.py:4: error: Overloaded function signatures 1 and 2 overlap with incompatible return types
tt.py:19: error: Revealed type is 'Union[builtins.int, builtins.None]'
tt.py:20: error: Revealed type is 'Any'
tt.py:21: error: Revealed type is 'Any'

I also tried None as the argument to the 2nd overload. That got me closer:

tt.py:19: error: Revealed type is 'Union[builtins.int, builtins.None]'
tt.py:20: error: Revealed type is 'Any'
tt.py:21: error: Revealed type is 'builtins.str'

I would love to be able to specify this as it would make a lot of our code easier to type. Maybe there's a way to do it with a TypeVar but I couldn't figure it out.

[gvanrossum]

And for anyone thinking that maybe type variables provide the solution, the 'str' type in the input is unrelated to the 'str' type in the output -- the real-world code from which this was derived has bytes -> str.

Also this reminds me of #1436.

[refi64]

IIRC you can do:

OptStr = TypeVar('OptStr', str, Optional[str])

# You might be able to instead do:
OptStr = TypeVar('OptStr', bound=Optional[str])
# but I'm not sure...

def bar(foo: OptStr) -> OptStr:
    # ...

That's actually how AnyStr is defined in the stubs.

[gvanrossum]

Our messages crossed, there is no type correspondence between arg and return except when arg is None.

[euresti]

Oops. Looks like I simplified my example too much. I edited the original so that future readers don't have to read all the comments about how my example was wrong. Sorry!

[sixolet]

I don't see any reason not to relax the overloading rules here; it feels like they're overly restrictive in this case. Specifically, since str is a subtype of Optional[str] and the valid calls with a int argument are a proper subset of the valid calls with an Optional[int] argument, then we should allow this particular type of overlapping signature.

[ilevkivskyi]

I think this should be done like this:

@overload
def bar(x: None) -> None:
    ...
@overload
def bar(x: int) -> str:
    ...
def bar(x: Optional[int]) -> Optional[str]:
    if x is None:
        return None
    return str(x)

However, it does not work. I get this:

x: Optional[int]

reveal_type(bar(1)) # OK, builtins.str
reveal_type(bar(None)) # Any + a spurious error: "bar" does not return a value
reveal_type(bar(x)) # Any, no error reported on this line

It looks like we have two independent bugs here. The spurious error "bar" does not return a value seems to be specific for None, so let us consider arbitrary types:

class A: ...
class B: ...
class C: ...

@overload
def foo(x: A) -> A: ...
@overload
def foo(x: B) -> C: ...
def foo(x: Union[A, B]) -> Union[A, C]:
    if isinstance(x, A):
        return x
    return C()

Here A plays the role of NoneType. With this code I get the following:

z: Union[A, B]

reveal_type(foo(A())) # OK, __main__.A
reveal_type(foo(B())) # OK, __main__.C
reveal_type(foo(z)) # Any, no error reported

The fact that Any is silently inferred instead of Union[A, C] seems to me like a bug in @overload implementation.

[gvanrossum]

Yes, you called it. Thanks @ilevkivskyi .

[sixolet]

I don't think we should infer the Union[A, C] return type for foo(z), and I think we should error in this case -- it seems like the implementation type should not factor in to the typechecking of an overload, otherwise you wouldn't be able to use the overloads to narrow specifically the types your function accepted. (You should remain able to make an implementation that takes (*args: Any, **kwargs: Any) and narrow the types you'll accept by defining an overload series).

Doing the type math to calculate that a type union of two arguments would result in a return type of a type union of two return types seems chancy and hard for the user to understand.

We should, on the other hand, allow you to define a third

@overload
def foo(x: Union[A, B]) -> Union[A, C]: ...

That would make everything typecheck the way we want it to.

[ilevkivskyi]

@sixolet

it seems like the implementation type should not factor in to the typechecking of an overload

Inferred Union[A, C] in my example has nothing to do with the implementation type, I just added some type in implementation to typecheck the body.

Doing the type math to calculate that a type union of two arguments would result in a return type of a type union of two return types seems chancy and hard for the user to understand.

There is a general "policy" now to not infer union types, but this is moving now towards inferring unions in simple cases, this was discussed with @JukkaL recently in the context of Instance joins. TypeScript for example always infers unions in the latter case, and people don't complain much about this. On the contrary, TypeScript does not support union math for overloads and people seem to complain a lot about this.

We can consider the option of the third overload you proposed (at least we should not prohibit it, there are already few open issues about order of overloads and overlapping types). Then later we could add some simple union math. Anyway, silently inferring Any is not a good idea in any case.

[JukkaL]

The fact that Any is silently inferred instead of Union[A, C] seems to me like a bug in @overload implementation.

Yes, this is a known bug. It looks like it's basically the same as #3256.