Function takes `Foo` subclass and wraps it in `Bar`, else returns type unaltered

Question

I have the following code:

from typing import TypeVar, Any, Generic


class Foo:
    ...

class Bar(Generic[FooT]):
    def __init__(self, foo: FooT):
        self._foo = foo


FooT = TypeVar('FooT', bound=Foo)
T = TypeVar('T')


def func(a: FooT | T) -> Bar[FooT] | T:
    if isinstance(a, Foo):
        return Bar(a)
    return a

def my_other_func(a: Foo) -> None:
    func(a)

• func takes either a Foo subclass and returns Bar which wraps that Foo object, or returns the input unaltered

I thought I could type it as

def func(a: FooT | T) -> Bar[FooT] | T:

But if I run mypy on this I get

main.py:18: error: Argument 1 to "Bar" has incompatible type "Foo"; expected "FooT"  [arg-type]
main.py:22: error: Argument 1 to "func" has incompatible type "Foo"; expected "Never"  [arg-type]
Found 2 errors in 1 file (checked 1 source file)

and I don't understand either.

How should I have typed it?

Answer 1

Use typing.overload:

from typing import TypeVar, Generic


class Foo:
    ...

FooT = TypeVar('FooT', bound=Foo)

class Bar(Generic[FooT]):
    def __init__(self, foo: FooT):
        self._foo = foo


T = TypeVar('T')


@overload
def func(a: FooT) -> Bar[FooT]:
    ...

@overload
def func(a: T) -> T:
    ...

def func(a):
    if isinstance(a, Foo):
        return Bar(a)
    return a

def my_other_func(a: Foo) -> None:
    func(a)

Essentially, Callable[[FooT|T], Bar[FooT]|T] is not the same as Callable[[FooT], Bar[FooT]] | Callable[[T], T]. typing.overload lets you define the latter instead of the former.