[Python] 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?

Continue reading...