add_slots tool breaks pickling of frozen dataclasses & proposed solution

[ariebovenberg]

Although not part of the official API, I find add_slots to be quite useful. However, the combination with frozen dataclasses and pickling causes problems:

@add_slots
@dataclass(frozen=True)
class ExampleDataclass:
    foo: str
    bar: int



assert ExampleDataclass.__slots__ == ("foo", "bar")

assert pickle.loads(
    pickle.dumps(ExampleDataclass("a", 1))
) == ExampleDataclass("a", 1)

gives the following error:

dataclasses.FrozenInstanceError: cannot assign to field 'foo'

A quick fix would be to add something like this:

def _dataclass_getstate(self):
    return [getattr(self, f.name) for f in fields(self)]


def _dataclass_setstate(self, state):
    for field, value in zip(fields(self), state):
        # use setattr because dataclass may be frozen
        object.__setattr__(self, field.name, value)


def add_slots(cls):
    ...  # existing add_slots code here...
    # optionally only do these steps if the dataclass is frozen
    cls.__getstate__ = _dataclass_getstate
    cls.__setstate__ = _dataclass_setstate
    return cls

edit: typos