pythonmypypython-typingpython-dataclasses

How to statically enforce frozen data classes in Python?


I'm trying to write an example where I'd like to use a frozen dataclass instance during type checking and swap it out with a normal dataclass to avoid paying the instantiation cost of frozen dataclasses.

The goal is to ensure that the instance is immutable with the type checker and use a regular dataclass during runtime. Here's the snippet:

from dataclasses import dataclass
from typing import TYPE_CHECKING
from functools import partial

if TYPE_CHECKING:
    frozen = partial(dataclass, frozen=True)
else:
    frozen = dataclass

@frozen
class Foo:
    x: int
    y: int

foo = Foo(1, 2) # mypy complains about the number of arguments

foo.x = 3 # instead, mypy should complain here

This works as expected during runtime, but running mypy raises this error. Pyright gives me the same error as well:

foo.py:49: error: Too many arguments for "Foo"  [call-arg]

In this snippet, the type checker can catch the mutation error:

@dataclass(frozen=True)
class Foo:
    x: int
    y: int

foo = Foo(1, 2)

foo.x = 3 # mypy correctly catches the error here

So, I'm guessing that the type checker doesn't like when I'm aliasing frozen = dataclass or frozen = partial(...). How do I annotate this properly so that the type checker understands that it's a dataclass instance and doesn't complain about mismatched argument count?

P.S: This is just an exercise. I know turning on dataclass(frozen=True) is way easier, and I shouldn't care about performance in such cases. I was inspired to try this after reading a blog post by Tin Tvrtković on making attr class instances frozen at compile time.


Solution

  • Use @dataclass_transform with frozen_default = True:

    if TYPE_CHECKING:
        T = TypeVar('T')
        
        @dataclass_transform(frozen_default = True)
        def frozen(cls: type[T]) -> type[T]:
            ...
    else:
        frozen = dataclass
    

    frozen_default was added in Python 3.12. However, since @dataclass_transform deliberately accepts all keyword arguments, 3.11 (exactly the same) and lower (using typing_extensions) will allow this just fine.

    This works with both Mypy and Pyright:

    reveal_type(Foo)  # mypy    => (x: int, y: int) -> Foo
                      # pyright => type[Foo]
    
    foo = Foo(1, 2)   # mypy + pyright => fine
    
    foo.x = 3         # mypy    => error: Property "x" defined in "Foo" is read-only
                      # pyright => error: Cannot assign member "x" for type "Foo"; "Foo" is frozen