jcrist / msgspec

A fast serialization and validation library, with builtin support for JSON, MessagePack, YAML, and TOML
https://jcristharif.com/msgspec/
BSD 3-Clause "New" or "Revised" License
2.44k stars 75 forks source link

"partially-frozen" structs? #716

Open inklesspen opened 3 months ago

inklesspen commented 3 months ago

Description

I've been converting attrs-based classes to structs in one of my projects, and I have a use case for a partially-frozen struct.

import datetime
import typing
import uuid

import msgspec

class User(msgspec.Struct, kw_only=True):
    id: uuid.UUID
    group_ids: tuple[uuid.UUID, ...] = msgspec.field(default_factory=tuple)
    join_date: datetime.datetime
    expiration_date: typing.Optional[datetime.datetime] = None
    notes: str

    def __setattr__(self, name: str, value: typing.Any) -> None:
        if name != "notes":
            raise AttributeError(f"The field {name!r} cannot be modified.")
        return super().__setattr__(name, value)

    def to_dict(self):
        return msgspec.structs.asdict(self)

if __name__ == "__main__":
    some_user = User(
        id=uuid.uuid4(),
        group_ids=(uuid.uuid4(), uuid.uuid4()),
        join_date=datetime.datetime(year=1969, month=7, day=20, hour=20, minute=17, second=40, tzinfo=datetime.timezone.utc),
        notes="Initial notes.",
    )
    some_dict = some_user.to_dict()
    assert some_dict["notes"] == "Initial notes."
    json_user = msgspec.json.encode(some_user)
    print(json_user)
    decoded_user = msgspec.json.decode(json_user, type=User)
    assert decoded_user == some_user
    some_user.notes += "\nSubsequent notes."
    some_dict = some_user.to_dict()
    assert some_dict["notes"] == "Initial notes.\nSubsequent notes."
    print(msgspec.json.encode(some_user))
    try:
        some_user.group_ids += (uuid.uuid4(),)
    except AttributeError:
        pass
    else:
        print("Was able to modify a frozen field!")

Basically, I would like to have all the fields except notes be frozen. Additionally, I don't want notes to be taken into account for __hash__ and __eq__ (or the ordering methods, if order is True), and I expect it should be excluded from the pattern-matching support (__match_args__) too. The non-frozen field should still be encoded/decoded, but it doesn't matter in terms of the "value" of the struct.

(Of course, I can use my recipe here, but I would also have to implement my own __hash__ and __eq__ methods, which I'm not looking forward to.)

I know msgspec.structs.force_setattr exists, but I'm not 100% clear on when it is more or less safe to use it. However I do know that using it alters the hash for the struct, which doesn't work with what I want here.

If I had to design an API for this, I think would have a keep_thawed keyword argument in msgspec.field, with a doc note that it does nothing unless the struct is frozen, and if frozen, it excludes that field from the various value-related dunder methods.

class FrozenUser(msgspec.Struct, frozen=True, kw_only=True):
    id: uuid.UUID
    group_ids: tuple[uuid.UUID, ...] = msgspec.field(default_factory=tuple)
    join_date: datetime.datetime
    expiration_date: typing.Optional[datetime.datetime] = None
    notes: str = msgspec.field(default="", keep_thawed=True)

    def to_dict(self):
        return msgspec.structs.asdict(self)