Generic type interplay with Union type

Question

I use a Generic class in python, whose instances are little but wrappers around instances of other types:

import random
from typing import Generic
from typing import TypeVar
from typing import Union

T = TypeVar("T", covariant=True)


class Wrapper(Generic[T]):
    def __init__(self, value: T) -> None:
        self._value = value

    @property
    def unwrap(self) -> T:
        return self._value


def test_union_wrapper() -> None:
    def wrapper_union() -> Wrapper[Union[str, int]]:
        if random.random() >= 0.5:
            return Wrapper("foo")
        else:
            return Wrapper(42)

    # mypy will give an error for this line
    w_u: Union[Wrapper[str], Wrapper[int]] = wrapper_union()

Running mypy on the above code will result in:

error: Incompatible types in assignment (expression has type "Wrapper[Union[str, int]]", variable has type "Union[Wrapper[str], Wrapper[int]]")

This might seem reasonable, because

Wrapper[Union[str, int]] ≮: Wrapper[str], and Wrapper[Union[str, int]] ≮: Wrapper[int]

and, as can be read in PEP483:

Union behaves covariantly in all its arguments. Indeed, as discussed above, Union[t1, t2, ...] is a subtype of Union[u1, u2, ...], if t1 is a subtype of u1, etc.

But I take issue with it, because I know that ∀w ∈ Wrapper[Union[str, int]], w ∈ Wrapper[str] or w ∈ Wrapper[int], hence, Wrapper[Union[str, int]] <: Union[Wrapper[str], Wrapper[int]], regardless. I would like to get mypy to recognize the same fact, but I don't know how.

There is even an example of such a recognition using the standard library. If I replace Wrapper with Type - another covariant Generic - in the code above, we get:

def test_union_type() -> None:
    def type_union() -> Type[Union[str, int]]:
        if random.random() >= 0.5:
            return str
        else:
            return int

    # mypy has no problem with this
    w_u: Union[Type[str], Type[int]] = type_union()

Here, mypy recognizes that the function return type, Type[Union[str, int]] is equivalent to the variable type, Union[Type[str], Type[int]].

I thus have two questions:

How do I tell a type checker that Wrapper behaves as Type does with regards to Union?
Does this behavior with regards to Union have a name? If we treat Wrapper and Unions as functions on types, we could say that they commute with each other, but not sure what the right term is in the context of type theory, or Python specifically.

user_5969611 · Answer

I found an unsatisfying solution to Q1, using a mypy pluging with the following type analyzer hook:

def commute_with_union_hook(ctx: AnalyzeTypeContext) -> Type:
    sym = ctx.api.lookup_qualified(ctx.type.name, ctx.type)
    node = sym.node
    if not ctx.type.args and len(ctx.type.args) == 1:
        return ctx.api.analyze_type_with_type_info(node, ctx.type.args, ctx.type)

    t: Type = ctx.type.args[0]
    t = ctx.api.anal_type(t)

    if not isinstance(t, UnionType):
        wrapper = ctx.api.analyze_type_with_type_info(node, ctx.type.args, ctx.type)
        return wrapper
    else:
        union_wrapper = UnionType.make_union(
            items=[Instance(typ=node, args=[item]) for item in t.items]
        )
        return union_wrapper

Feels overkill, but it might be the only way.

Generic type interplay with Union type

1 Answers