How to use TypeVar for input and output of multiple generic Protocols in python?

Question

I want to use multiple generic protocols and ensure they're compatible:

from typing import TypeVar, Protocol, Generic
from dataclasses import dataclass

# checking fails as below and with contravariant=True or covariant=True:
A = TypeVar("A") 

class C(Protocol[A]):
    def f(self, a: A) -> None: pass

class D(Protocol[A]):
    def g(self) -> A: pass

# Just demonstrates my use case; doesn't have errors:
@dataclass
class CompatibleThings(Generic[A]):
    c: C[A]
    d: D[A]

Mypy gives the following error:

Invariant type variable 'A' used in protocol where contravariant one is expected
Invariant type variable 'A' used in protocol where covariant one is expected

I know this can be done by making C and D generic ABC classes, but I want to use protocols.

Answer 1

The short explanation is that your approach breaks subtype transitivity; see this section of PEP 544 for more information. It gives a pretty clear explanation of why your D protocol (and, implicitly, your C protocol) run into this problem, and why it requires different types of variance for each to solve it. You can also look on Wikipedia for info on type variance.

Here's the workaround: use covariant and contravariant protocols, but make your generic dataclass invariant. The big hurdle here is inheritance, which you have to handle in order to use Protocols, but is kind of tangential to your goal. I'm going to switch naming here to highlight the inheritance at play, which is what this is all about:

A = TypeVar("A") # Invariant type
A_cov = TypeVar("A_cov", covariant=True) # Covariant type
A_contra = TypeVar("A_contra", contravariant=True) # Contravariant type

# Give Intake its contravariance
class Intake(Protocol[A_contra]):
    def f(self, a: A_contra) -> None: pass

# Give Output its covariance
class Output(Protocol[A_cov]):
    def g(self) -> A_cov: pass

# Just tell IntakeOutput that the type needs to be the same
# Since a is invariant, it doesn't care that
# Intake and Output require contra / covariance
@dataclass
class IntakeOutput(Generic[A]):
    intake: Intake[A]
    output: Output[A]

You can see that this works with the following tests:

class Animal:
    ...
    
class Cat(Animal):
    ...
    
class Dog(Animal):
    ...
    
class IntakeCat:
    def f(self, a: Cat) -> None: pass

class IntakeDog:
    def f(self, a: Dog) -> None: pass

class OutputCat:
    def g(self) -> Cat: pass

class OutputDog:
    def g(self) -> Dog: pass

compat_cat: IntakeOutput[Cat] = IntakeOutput(IntakeCat(), OutputCat())
compat_dog: IntakeOutput[Dog] = IntakeOutput(IntakeDog(), OutputDog())

# This is gonna error in mypy
compat_fail: IntakeOutput[Dog] = IntakeOutput(IntakeDog(), OutputCat())

which gives the following error:

main.py:48: error: Argument 2 to "IntakeOutput" has incompatible type "OutputCat"; expected "Output[Dog]"
main.py:48: note: Following member(s) of "OutputCat" have conflicts:
main.py:48: note:     Expected:
main.py:48: note:         def g(self) -> Dog
main.py:48: note:     Got:
main.py:48: note:         def g(self) -> Cat

So what's the catch? What are you giving up? Namely, inheritance in IntakeOutput. Here's what you can't do:

class IntakeAnimal:
    def f(self, a: Animal) -> None: pass

class OutputAnimal:
    def g(self) -> Animal: pass

# Ok, as expected
ok1: IntakeOutput[Animal] = IntakeOutput(IntakeAnimal(), OutputAnimal())

# Ok, because Output is covariant
ok2: IntakeOutput[Animal] = IntakeOutput(IntakeAnimal(), OutputDog())

# Both fail, because Intake is contravariant
fails1: IntakeOutput[Animal] = IntakeOutput(IntakeDog(), OutputDog())
fails2: IntakeOutput[Animal] = IntakeOutput(IntakeDog(), OutputAnimal())

# Ok, because Intake is contravariant
ok3: IntakeOutput[Dog] = IntakeOutput(IntakeAnimal(), OutputDog())

# This fails, because Output is covariant
fails3: IntakeOutput[Dog] = IntakeOutput(IntakeAnimal(), OutputAnimal())
fails4: IntakeOutput[Dog] = IntakeOutput(IntakeDog(), OutputAnimal())

So. There it is. You can play around with this more here.