[SOLVED] Python: Implement base class for different types

Python: Implement base class for different types

I wonder if there is a way to implement subclasses of a base class for different types. Each subclass should have individual input and output types while providing same behaviour as the base class.

Background: I want to process voltage and temperature samples. 100 voltage samples form a VoltageDataset. 100 temperature samples form a TemperatureDataset. Multiple VoltageDatasets form a VoltageDataCluster. Same for temperature. The processing of Datasets depends on their physical quantity. To ensure that voltage related processing can't be applied to temperature samples I'd like to add type hints.

So I'd would be nice if there is a way to define that VoltageDataClustes method append_dataset allows VoltageDataset as input type only. Same for temperature.

Is there a way to implement this behaviour without copy&pasting?

# base class
class DataCluster:
    def __init__(self, name):
        self.name = name
        self.datasets = list()

    def append_dataset(self, dataset: Dataset) -> None:
        self.datasets.append(dataset)

# subclass that should allow VoltageDataset input only.
class VoltageDataCluster(DataCluster):
    pass

# subclass that should allow TemperatureDataset input only.
class TemperatureDataCluster(DataCluster):
    pass

Thanks! Niklas

Solution

You could use pydantic generic models.

from typing import Generic, TypeVar, List
from pydantic.generics import GenericModel

DataT = TypeVar('DataT')

class DataCluster(GenericModel, Generic[DataT]):
    name: str
    datasets: List[DataT] = []

    def append_dataset(self, dataset: DataT) -> None:
        self.datasets.append(dataset)

voltage_cluster = DataCluster[VoltageDataset](name="name")
voltage_cluster.append_dataset(some_voltage_dataset)