dance_mypy_dance.py

from typing import TypeVar, Generic, Literal, Final, TypedDict, Protocol, Any

T = TypeVar("T")
U = TypeVar("U")
T_contra = TypeVar("T_contra", contravariant=True)
U_co = TypeVar("U_co", covariant=True)


class Metric(Protocol[T_contra, U_co]):
    def add(self, value: T_contra) -> None:
        ...
    
    def flush(self) -> U_co:
        ...
        
class CounterMetric(Metric[float, float]):
    value: float
    
    def __init__(self):
        self.value = 0.0

    def add(self, value: float) -> None:
        self.value += value
    
    def flush(self) -> float:
        return self.value
        
class StrListMetric(Metric[str, list[str]]):
    value: list[str]
    
    def __init__(self):
        self.value = []
    
    def add(self, value: str) -> None:
        self.value += [value]
    
    def flush(self) -> list[str]:
        return self.value


class NamedMetric(Generic[T, U]):
    name: str
    metric_constructor: type[Metric[T, U]]
    
    def __init__(self, name: str, metric_constructor: type[Metric[T, U]]) -> None:
        self.name = name
        self.metric_constructor = metric_constructor

class Aggregator:
    buckets: dict[Any, Any] # typing guaranteed by logic
    
    def add(self, m: NamedMetric[T, U], key: str, v: T) -> None:
        if m.name not in self.buckets:
            metric = self.buckets[(m.name, key)] = m.metric_constructor()
        else:
            metric = self.buckets[(m.name, key)]
        metric.add(v)

A_METRIC = NamedMetric("a", CounterMetric)

a = Aggregator()
a.add(A_METRIC, "foo", 123) # type checks
a.add(A_METRIC, "foo", "123") # errors, but a terrible one: Cannot infer type argument 1 of "add" of "Aggregator"  [misc]