How do I specify "related" types in a container class?

I have various model classes, each working on its own data class. Like so:

from dataclasses import dataclass
from typing import Any, Protocol, final

@dataclass
class Data(Protocol):
    val_generic: int

@dataclass
class DataA(Data):
    val_generic: int = 1
    val_a: int = 2

@dataclass
class DataB(Data):
    val_generic: int = 4
    val_b: int = 1

class ModelA:
    def update(self, data: DataA) -> None:
        data.val_a = data.val_a + data.val_generic

class ModelB:
    def update(self, data: DataB) -> None:
        data.val_b = data.val_b + data.val_generic

Now I want a container class doing this job for me:

@final
class Container:
    def __init__(self, data: Data, model: Any):
        self.data = data
        self.model = model

    def update(self):
        self.model.update(self.data)

model_a = ModelA()
data_a = DataA()
container1 = Container(data_a, model_a)
container1.update() # OK

I fail to give a valid type annotation to the model-argument. In an ideal world, the type checker realizes that the code above works while this will raise an error:

data_b = DataB()
container2 = Container(data_b, model_a)
container2.update() #AttributeError

Can it be done in a generic fashion? Things are not as easy as in this MRE. I tried to define a Model-protocol but failed to write an abstract update-method since the signature of ModelA.update differs from ModelB.update.

But even if that works out, how do I tell the type-checker that the argument-types in Container.__init__ have to be related in order to make Container.update work?

Solution

If you can't change the definitions of ModelA and ModelB, a nominal approach won't work since your model and data types aren't related to each other.

Instead, define a generic protocol for models that support updating with a certain Data subtype:

class Model[T: Data](Protocol):
    def update(self, data: T) -> None: ...

Then, if you make Container generic on the type of data being stored, you can ensure any arguments passed as model support updates from that data.

@final
class Container[T: Data]:
    def __init__(self, data: T, model: Model[T]):
        self.data = data
        self.model = model

    def update(self):
        self.model.update(self.data)

If you ever try to instantiate a Container with mismatched arguments, you should get a warning right at the constructor. Here's a small test suite to confirm:

model_a = ModelA()
data_a = DataA()
model_b = ModelB()
data_b = DataB()

containerAA = Container(data_a, model_a)
containerBB = Container(data_b, model_b)
containerAB = Container(data_a, model_b) # Bad
ContainerBA = Container(data_b, model_a) # Bad

Pyright's output on this is pretty verbose (it usually is when you're working with protocols), but it does catch all the problems.

$ pyright container_types.py
c:\path\to\folder\container_types.py
  c:\path\to\folder\container_types.py:45:33 - error: Argument of type "ModelB" cannot be assigned to parameter "model" of type "Model[T@Container]" in function "__init__"
    "ModelB" is incompatible with protocol "Model[DataA]"
      "update" is an incompatible type
        Type "(data: DataB) -> None" is not assignable to type "(data: T@Model) -> None"
          Parameter 1: type "T@Model" is incompatible with type "DataB"
            "DataA" is not assignable to "DataB" (reportArgumentType)
  c:\path\to\folder\container_types.py:46:33 - error: Argument of type "ModelA" cannot be assigned to parameter "model" of type "Model[T@Container]" in function "__init__"
    "ModelA" is incompatible with protocol "Model[DataB]"
      "update" is an incompatible type
        Type "(data: DataA) -> None" is not assignable to type "(data: T@Model) -> None"
          Parameter 1: type "T@Model" is incompatible with type "DataA"
            "DataB" is not assignable to "DataA" (reportArgumentType)
2 errors, 0 warnings, 0 informations