Generic type interplay with Union type

[harahu]

I have already posted a lengthy version of this question on stackoverflow, but I discovered this forum and thought it would be a just as good place to ask. Extended version: https://stackoverflow.com/questions/69209879/generic-type-interplay-with-union-type

TLDR:

Does there exist a way to specify that a type commutes with Union? In other words, how do I get this example past a type checker?

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()

[sobolevn]

Typechecks cannot be sure if Union[Wrapper[A], Wrapper[B]] is the same as Wrapper[Union[A, B]]. Why?
Imagine that Wrapper has some .add(self, other: T) method that can add a value to the wrapper.

In case of Union[Wrapper[A], Wrapper[B]] you can only use .add(A()) for the first case and only use .add(B()) for the second case. While with Wrapper[Union[A, B]] you can add A, B, and A | B (depending on the variance of T).

So, why won't you just stick to the first or the second form? 🙂

[harahu]

The function above is just one example. Say I want to do something like:

a: Union[str, int] = ...
b = Wrapper(a)

c: Union[Foo, Bar] = ...
d = Wrapper(c)

Here I cannot control the form of the type of b or d, without lots of verbose code., so it then makes sense to adopt the Wrapper[Union[...]] form.

Now, I'll struggle maintaining this later on if I want to do the following:

e = b if ... else d

Suddenly I have the type Union[Wrapper[Union[...]], Wrapper[Union[...]]]. For complex problems this can grow quite bad, and I want a way of collapsing these expressions to the simplest form, that is, in this case:

e: Union[Wrapper[str], Wrapper[int], Wrapper[Foo], Wrapper[Bar]]

[harahu]

As for your point regarding .add(self, other: T): I understand why this would be problematic, and I don't expect the type checker to understand that there is no such method on its own, without me claiming it somehow.

I guess I am looking for a pattern similar to how you can tell the type checker that a Generic is covariant or contravariant with regards to a TypeVar. It seems to me like it might not exist.

[harahu]

Was able to achieve the effect I was going for using a custom mypy plugin with the following type analyze hook:

def commute_with_union_hook(ctx: AnalyzeTypeContext) -> Type:
    if ctx.type.args and len(ctx.type.args) == 1:
        t: Type = ctx.type.args[0]
        t = ctx.api.anal_type(t)

        sym = ctx.api.lookup_qualified(ctx.type.name, ctx.type)
        node = sym.node

        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

    return ctx.type

This feels quite crude, but it works for the examples I had in mind.

Heavily inspired by the mypy.types.TypeType type.