Smarter Type Narrowing With Intersection + Union

[zevbo]

Consider the following python code:

from enum import Enum
from typing import Generic, Literal, TypeVar

class E(Enum):
    A = 0 
    B = 1

T = TypeVar("T", Literal[E.A], Literal[E.B])

class A(Generic[T]): 
    pass 

class B(A[Literal[E.A]]): 
    pass 

class C(A[Literal[E.B]]):
    pass 

TheUnion = B | C

Now consider the following function:

def a(b: A[E.A] & TheUnion) -> B:
    return b.

Here, there is enough information to narrow the type of the argument b to of type B, as there only one of the classes in the Union that can be of type A[E.A]. C nor any subclass of C can have type A[E.A], as a result of C being of type A[E.B].

However, basedmypy is unable to make this narrowing, and instead says:
Incompatible return value type (got "A[E.A] & B | C", expected "B") [return-value]

I would love if basedmypy could correctly make this inference.

[KotlinIsland]

I love the idea, but unfortunately I don't believe the suggestion is entirely accurate, consider the case in which a third subclass of A (with E.A) is declared, the class also extends C:

class D(C, A[Literal[E.A]]):
    pass

This class is not a B, but is an A[E.A] & TheUnion.

Perhaps if C was declared as final, then b would be a B. A similar issue is being tracked in #451.

Edit: sealed would have the exact same issue as in the OP

Another feature that would make this scenario work would be sealed classes (#204), and would simplify the situation by making the union redundant:

from enum import Enum
from typing import Generic, Literal, TypeVar

class E(Enum):
    A = 0 
    B = 1

T = TypeVar("T", Literal[E.A], Literal[E.B])

@sealed
class A(Generic[T]): 
    pass 

class B(A[Literal[E.A]]): 
    pass 

class C(A[Literal[E.B]]):
    pass 

def a(b: A[E.A]) -> B:
    return b

[zevbo]

"I love the idea, but unfortunately I don't believe the suggestion is entirely accurate, consider the case in which a third subclass of A (with E.A) is declared, the class also extends C:"

I had the same thought. But when I do the following:

0: class D(C, A[Literal[E.A]]):
1:     pass
2: doesnt_work: A[Literal[E.A]] = D()

Line 2 get: Incompatible types in assignment (expression has type "D", variable has type "A[E.A]") [assignment].
With that in mind, it seems like the narrowing is correct? I'm not 100% sure though.

And yes, I certainly agree with you on @closed--and that certainly gets past any concerns. But it seemed to me like that feature had been scrapped? But I'd love to hear if it's not dead in the water.

[KotlinIsland]

Oh yeah true, A[E.A] & A[E.B] is impossible. I think mypy should have reported an error on the class def of D.

I agree that basedmypy should narrow here. I'll promote it to an issue.

But it seemed to me like that feature had been scrapped? But I'd love to hear if it's not dead in the water.

The way it works in python is if someone can be bothered to make it, it will get implemented.

[KotlinIsland]

Can I ask, what is your 'mypy in real life' moment that unearthed this?

[KotlinIsland]

After further investigation by the team, we have concluded that this example is an instance of #432.

with A[E.A] & (B | C) == (A[E.A] & B) | (A[E.A] & C). The A[E.A] & C should become Never and disappear, leaving only the A[E.A] & B, which simplifies to B.

I've added this example to #432. Please leave an upvote.

Thank you very much for using basedmypy and starting this discussion. 🚀

[KotlinIsland]

I think mypy should have reported an error on the class def of D.

• (🎁) Report when generic base class specified multiple times with incompatible types python/mypy#15348

[zevbo]

That's what I thought too

[zevbo]

Can I ask, what is your 'mypy in real life' moment that unearthed this?

Yeah sure. So I have a set of coresponding request and response types that look vaguely like:

class BaseRequest: ... 
class BaseResponse: ...

class ARequest(BaseReqeuest): ...
class AResponse(BaseResponse): ...

class BRequest(BaseReqeuest): ...
class BResponse(BaseResponse): ...

class CRequest(BaseReqeuest): ...
class CResponse(BaseResponse): ...

There are a number of things I want to do, that I need to generalize to each "kind" of request-response pair. For example, I might want a container class that looks like:

RequestT = TypeVar("RequestT", bound=BaseRequest)
ResponseT = TypeVar("ResponseT", bound=BaseResponse)
@dataclass
class RequestResponseContainer(Generic[RequestT, ResponseT]):
    request: RequestT
    response: ResponseT

Now, if I want to associated each request and response, I really have to make subclasses for each one. Plus, I'll need to subtype so that if I give you a RequestResponseContainer[ARequest, AResponse], mypy won't be able to determine the actual class that request and response are.

So what I would really love to be able to do is the following [please excuse me not using Literal everywhere correctly]:

class Kind(Enum):
    A = 0
    B = 1
    C = 2

KindT = TypeVar("KindT", Literal[Kind.A], Literal[Kind.B], Literal[ThingsKindC])

class BaseRequest(Generic[KindT]): ... 
class BaseResponse(Generic[KindT]): ...

class ARequest(BaseReqeuest[Literal[Kind.A]): ...
class AResponse(BaseResponse[Literal[Kind.A]]): ...

class BRequest(BaseReqeuest[Literal[Kind.B]])): ...
class BResponse(BaseResponse[Literal[Kind.B]])): ...

class CRequest(BaseReqeuest[Literal[Kind.C]])): ...
class CResponse(BaseResponse[Literal[Kind.C]])): ...

RequestType = BaseRequest[KindT] & (Kind.A | Kind.B | Kind.C)
ResponseType = BaseResponse[KindT] & (Kind.A | Kind.B | Kind.C)

@dataclass
class RequestResponseContainer(Generic[KindT]):
    request: RequestType[KindT]
    response: ResponseType[KindT]

Thus, the request and response are already assocaited, and if I give you a RequestResponseContainer[Kind.A], you know that request is an ARequest and response is an AResponse. Obviously this would be better (and kind of more conceptually correct) with @sealed as you wouldn't need the union--but that seemed like a path of more resistance.

And in my project I have quite a few places where I want this safety on my request/response abstraction, so it would actually be super helpful.

[KotlinIsland]

I don't 100% understand, but could something like the following work???

from __future__ import annotations

from dataclasses import dataclass
from typing import Generic, TypeVar

class Kind: ...
class A(Kind): ...
class B(Kind): ...
class C(Kind): ...

class BaseRequest: ...
class BaseResponse: ...

class ARequest(BaseRequest, A): ...
class AResponse(BaseResponse, A): ...

class BRequest(BaseRequest, B): ...
class BResponse(BaseResponse, B): ...

class CRequest(BaseRequest, C): ...
class CResponse(BaseResponse, C): ...

KindT = TypeVar("KindT", bound=Kind)

@dataclass
class RequestResponseContainer(Generic[KindT]):
    request: BaseRequest & KindT
    response: BaseResponse & KindT


x = RequestResponseContainer[B](BRequest(), BResponse())