Search code examples
pythonmypypython-dataclasses

Annotate dataclass class variable with type value


We have a number of dataclasses representing various results with common ancestor Result. Each result then provides its data using its own subclass of ResultData. But we have trouble to annotate the case properly.

We came up with following solution:

from dataclasses import dataclass
from typing import ClassVar, Generic, Optional, Sequence, Type, TypeVar


class ResultData:
    ...


T = TypeVar('T', bound=ResultData)


@dataclass
class Result(Generic[T]):
    _data_cls: ClassVar[Type[T]]
    data: Sequence[T]

    @classmethod
    def parse(cls, ...) -> T:
        self = cls()
        self.data = [self._data_cls.parse(...)]
        return self

class FooResultData(ResultData):
    ...

class FooResult(Result):
    _data_cls = FooResultData

but it stopped working lately with mypy error ClassVar cannot contain type variables [misc]. It is also against PEP 526, see https://www.python.org/dev/peps/pep-0526/#class-and-instance-variable-annotations, which we missed earlier.

Is there a way to annotate this case properly?


Solution

  • At the end I just replaced the variable in _data_cls annotation with the base class and fixed the annotation of subclasses as noted by @rv.kvetch in his answer.

    The downside is the need to define the result class twice in every subclass, but in my opinion it is more legible than extracting the class in property.

    The complete solution:

    from dataclasses import dataclass
    from typing import ClassVar, Generic, Optional, Sequence, Type, TypeVar
    
    
    class ResultData:
        ...
    
    
    T = TypeVar('T', bound=ResultData)
    
    
    @dataclass
    class Result(Generic[T]):
        _data_cls: ClassVar[Type[ResultData]]  # Fixed annotation here
        data: Sequence[T]
    
        @classmethod
        def parse(cls, ...) -> T:
            self = cls()
            self.data = [self._data_cls.parse(...)]
            return self
    
    class FooResultData(ResultData):
        ...
    
    class FooResult(Result[FooResultData]):  # Fixed annotation here
        _data_cls = FooResultData