efro.dataclassio

29class Codec(Enum):
30    """Specifies expected data format exported to or imported from."""
31
32    # Use only types that will translate cleanly to/from json: lists,
33    # dicts with str keys, bools, ints, floats, and None.
34    JSON = 'json'
35
36    # Mostly like JSON but passes bytes and datetime objects through
37    # as-is instead of converting them to json-friendly types.
38    FIRESTORE = 'firestore'

 61class DataclassFieldLookup(Generic[T]):
 62    """Get info about nested dataclass fields in type-safe way."""
 63
 64    def __init__(self, cls: type[T]) -> None:
 65        self.cls = cls
 66
 67    def path(self, callback: Callable[[T], Any]) -> str:
 68        """Look up a path on child dataclass fields.
 69
 70        example:
 71          DataclassFieldLookup(MyType).path(lambda obj: obj.foo.bar)
 72
 73        The above example will return the string 'foo.bar' or something
 74        like 'f.b' if the dataclasses have custom storage names set.
 75        It will also be static-type-checked, triggering an error if
 76        MyType.foo.bar is not a valid path. Note, however, that the
 77        callback technically allows any return value but only nested
 78        dataclasses and their fields will succeed.
 79        """
 80
 81        # We tell the type system that we are returning an instance
 82        # of our class, which allows it to perform type checking on
 83        # member lookups. In reality, however, we are providing a
 84        # special object which captures path lookups, so we can build
 85        # a string from them.
 86        if not TYPE_CHECKING:
 87            out = callback(_PathCapture(self.cls))
 88            if not isinstance(out, _PathCapture):
 89                raise TypeError(
 90                    f'Expected a valid path under'
 91                    f' the provided object; got a {type(out)}.'
 92                )
 93            return out.path
 94        return ''
 95
 96    def paths(self, callback: Callable[[T], list[Any]]) -> list[str]:
 97        """Look up multiple paths on child dataclass fields.
 98
 99        Functionality is identical to path() but for multiple paths at once.
100
101        example:
102          DataclassFieldLookup(MyType).paths(lambda obj: [obj.foo, obj.bar])
103        """
104        outvals: list[str] = []
105        if not TYPE_CHECKING:
106            outs = callback(_PathCapture(self.cls))
107            assert isinstance(outs, list)
108            for out in outs:
109                if not isinstance(out, _PathCapture):
110                    raise TypeError(
111                        f'Expected a valid path under'
112                        f' the provided object; got a {type(out)}.'
113                    )
114                outvals.append(out.path)
115        return outvals

138class IOAttrs:
139    """For specifying io behavior in annotations.
140
141    'storagename', if passed, is the name used when storing to json/etc.
142    'store_default' can be set to False to avoid writing values when equal
143        to the default value. Note that this requires the dataclass field
144        to define a default or default_factory or for its IOAttrs to
145        define a soft_default value.
146    'whole_days', if True, requires datetime values to be exactly on day
147        boundaries (see efro.util.utc_today()).
148    'whole_hours', if True, requires datetime values to lie exactly on hour
149        boundaries (see efro.util.utc_this_hour()).
150    'whole_minutes', if True, requires datetime values to lie exactly on minute
151        boundaries (see efro.util.utc_this_minute()).
152    'soft_default', if passed, injects a default value into dataclass
153        instantiation when the field is not present in the input data.
154        This allows dataclasses to add new non-optional fields while
155        gracefully 'upgrading' old data. Note that when a soft_default is
156        present it will take precedence over field defaults when determining
157        whether to store a value for a field with store_default=False
158        (since the soft_default value is what we'll get when reading that
159        same data back in when the field is omitted).
160    'soft_default_factory' is similar to 'default_factory' in dataclass
161        fields; it should be used instead of 'soft_default' for mutable types
162        such as lists to prevent a single default object from unintentionally
163        changing over time.
164    'enum_fallback', if provided, specifies an enum value to be substituted
165        in the case of unrecognized enum values.
166    """
167
168    # A sentinel object to detect if a parameter is supplied or not.  Use
169    # a class to give it a better repr.
170    class _MissingType:
171        pass
172
173    MISSING = _MissingType()
174
175    storagename: str | None = None
176    store_default: bool = True
177    whole_days: bool = False
178    whole_hours: bool = False
179    whole_minutes: bool = False
180    soft_default: Any = MISSING
181    soft_default_factory: Callable[[], Any] | _MissingType = MISSING
182    enum_fallback: Enum | None = None
183
184    def __init__(
185        self,
186        storagename: str | None = storagename,
187        *,
188        store_default: bool = store_default,
189        whole_days: bool = whole_days,
190        whole_hours: bool = whole_hours,
191        whole_minutes: bool = whole_minutes,
192        soft_default: Any = MISSING,
193        soft_default_factory: Callable[[], Any] | _MissingType = MISSING,
194        enum_fallback: Enum | None = None,
195    ):
196        # Only store values that differ from class defaults to keep
197        # our instances nice and lean.
198        cls = type(self)
199        if storagename != cls.storagename:
200            self.storagename = storagename
201        if store_default != cls.store_default:
202            self.store_default = store_default
203        if whole_days != cls.whole_days:
204            self.whole_days = whole_days
205        if whole_hours != cls.whole_hours:
206            self.whole_hours = whole_hours
207        if whole_minutes != cls.whole_minutes:
208            self.whole_minutes = whole_minutes
209        if soft_default is not cls.soft_default:
210            # Do what dataclasses does with its default types and
211            # tell the user to use factory for mutable ones.
212            if isinstance(soft_default, (list, dict, set)):
213                raise ValueError(
214                    f'mutable {type(soft_default)} is not allowed'
215                    f' for soft_default; use soft_default_factory.'
216                )
217            self.soft_default = soft_default
218        if soft_default_factory is not cls.soft_default_factory:
219            self.soft_default_factory = soft_default_factory
220            if self.soft_default is not cls.soft_default:
221                raise ValueError(
222                    'Cannot set both soft_default and soft_default_factory'
223                )
224        if enum_fallback is not cls.enum_fallback:
225            self.enum_fallback = enum_fallback
226
227    def validate_for_field(self, cls: type, field: dataclasses.Field) -> None:
228        """Ensure the IOAttrs instance is ok to use with the provided field."""
229
230        # Turning off store_default requires the field to have either
231        # a default or a default_factory or for us to have soft equivalents.
232
233        if not self.store_default:
234            field_default_factory: Any = field.default_factory
235            if (
236                field_default_factory is dataclasses.MISSING
237                and field.default is dataclasses.MISSING
238                and self.soft_default is self.MISSING
239                and self.soft_default_factory is self.MISSING
240            ):
241                raise TypeError(
242                    f'Field {field.name} of {cls} has'
243                    f' neither a default nor a default_factory'
244                    f' and IOAttrs contains neither a soft_default'
245                    f' nor a soft_default_factory;'
246                    f' store_default=False cannot be set for it.'
247                )
248
249    def validate_datetime(
250        self, value: datetime.datetime, fieldpath: str
251    ) -> None:
252        """Ensure a datetime value meets our value requirements."""
253        if self.whole_days:
254            if any(
255                x != 0
256                for x in (
257                    value.hour,
258                    value.minute,
259                    value.second,
260                    value.microsecond,
261                )
262            ):
263                raise ValueError(
264                    f'Value {value} at {fieldpath} is not a whole day.'
265                )
266        elif self.whole_hours:
267            if any(
268                x != 0 for x in (value.minute, value.second, value.microsecond)
269            ):
270                raise ValueError(
271                    f'Value {value} at {fieldpath}' f' is not a whole hour.'
272                )
273        elif self.whole_minutes:
274            if any(x != 0 for x in (value.second, value.microsecond)):
275                raise ValueError(
276                    f'Value {value} at {fieldpath}' f' is not a whole minute.'
277                )

41class IOExtendedData:
42    """A class that data types can inherit from for extra functionality."""
43
44    def will_output(self) -> None:
45        """Called before data is sent to an outputter.
46
47        Can be overridden to validate or filter data before
48        sending it on its way.
49        """
50
51    @classmethod
52    def will_input(cls, data: dict) -> None:
53        """Called on raw data before a class instance is created from it.
54
55        Can be overridden to migrate old data formats to new, etc.
56        """
57
58    def did_input(self) -> None:
59        """Called on a class instance after created from data.
60
61        Can be useful to correct values from the db, etc. in the
62        type-safe form.
63        """
64
65    # pylint: disable=useless-return
66
67    @classmethod
68    def handle_input_error(cls, exc: Exception) -> Self | None:
69        """Called when an error occurs during input decoding.
70
71        This allows a type to optionally return substitute data
72        to be used in place of the failed decode. If it returns
73        None, the original exception is re-raised.
74
75        It is generally a bad idea to apply catch-alls such as this,
76        as it can lead to silent data loss. This should only be used
77        in specific cases such as user settings where an occasional
78        reset is harmless and is preferable to keeping all contained
79        enums and other values backward compatible indefinitely.
80        """
81        del exc  # Unused.
82
83        # By default we let things fail.
84        return None
85
86    # pylint: enable=useless-return

 92class IOMultiType(Generic[EnumT]):
 93    """A base class for types that can map to multiple dataclass types.
 94
 95    This enables usage of high level base classes (for example
 96    a 'Message' type) in annotations, with dataclassio automatically
 97    serializing & deserializing dataclass subclasses based on their
 98    type ('MessagePing', 'MessageChat', etc.)
 99
100    Standard usage involves creating a class which inherits from this
101    one which acts as a 'registry', and then creating dataclass classes
102    inheriting from that registry class. Dataclassio will then do the
103    right thing when that registry class is used in type annotations.
104
105    See tests/test_efro/test_dataclassio.py for examples.
106    """
107
108    @classmethod
109    def get_type(cls, type_id: EnumT) -> type[Self]:
110        """Return a specific subclass given a type-id."""
111        raise NotImplementedError()
112
113    @classmethod
114    def get_type_id(cls) -> EnumT:
115        """Return the type-id for this subclass."""
116        raise NotImplementedError()
117
118    @classmethod
119    def get_type_id_type(cls) -> type[EnumT]:
120        """Return the Enum type this class uses as its type-id."""
121        out: type[EnumT] = cls.__orig_bases__[0].__args__[0]  # type: ignore
122        assert issubclass(out, Enum)
123        return out
124
125    @classmethod
126    def get_type_id_storage_name(cls) -> str:
127        """Return the key used to store type id in serialized data.
128
129        The default is an obscure value so that it does not conflict
130        with members of individual type attrs, but in some cases one
131        might prefer to serialize it to something simpler like 'type'
132        by overriding this call. One just needs to make sure that no
133        encompassed types serialize anything to 'type' themself.
134        """
135        return '_dciotype'

28class JsonStyle(Enum):
29    """Different style types for json."""
30
31    # Single line, no spaces, no sorting. Not deterministic.
32    # Use this where speed is more important than determinism.
33    FAST = 'fast'
34
35    # Single line, no spaces, sorted keys. Deterministic.
36    # Use this when output may be hashed or compared for equality.
37    SORTED = 'sorted'
38
39    # Multiple lines, spaces, sorted keys. Deterministic.
40    # Use this for pretty human readable output.
41    PRETTY = 'pretty'

100def dataclass_from_dict(
101    cls: type[T],
102    values: dict,
103    *,
104    codec: Codec = Codec.JSON,
105    coerce_to_float: bool = True,
106    allow_unknown_attrs: bool = True,
107    discard_unknown_attrs: bool = False,
108) -> T:
109    """Given a dict, return a dataclass of a given type.
110
111    The dict must be formatted to match the specified codec (generally
112    json-friendly object types). This means that sequence values such as
113    tuples or sets should be passed as lists, enums should be passed as
114    their associated values, nested dataclasses should be passed as dicts,
115    etc.
116
117    All values are checked to ensure their types/values are valid.
118
119    Data for attributes of type Any will be checked to ensure they match
120    types supported directly by json. This does not include types such
121    as tuples which are implicitly translated by Python's json module
122    (as this would break the ability to do a lossless round-trip with
123    data).
124
125    If coerce_to_float is True, int values passed for float typed fields
126    will be converted to float values. Otherwise, a TypeError is raised.
127
128    If `allow_unknown_attrs` is False, AttributeErrors will be raised for
129    attributes present in the dict but not on the data class. Otherwise,
130    they will be preserved as part of the instance and included if it is
131    exported back to a dict, unless `discard_unknown_attrs` is True, in
132    which case they will simply be discarded.
133    """
134    val = _Inputter(
135        cls,
136        codec=codec,
137        coerce_to_float=coerce_to_float,
138        allow_unknown_attrs=allow_unknown_attrs,
139        discard_unknown_attrs=discard_unknown_attrs,
140    ).run(values)
141    assert isinstance(val, cls)
142    return val

145def dataclass_from_json(
146    cls: type[T],
147    json_str: str,
148    coerce_to_float: bool = True,
149    allow_unknown_attrs: bool = True,
150    discard_unknown_attrs: bool = False,
151) -> T:
152    """Return a dataclass instance given a json string.
153
154    Basically dataclass_from_dict(json.loads(...))
155    """
156
157    return dataclass_from_dict(
158        cls=cls,
159        values=json.loads(json_str),
160        coerce_to_float=coerce_to_float,
161        allow_unknown_attrs=allow_unknown_attrs,
162        discard_unknown_attrs=discard_unknown_attrs,
163    )

44def dataclass_to_dict(
45    obj: Any,
46    codec: Codec = Codec.JSON,
47    coerce_to_float: bool = True,
48    discard_extra_attrs: bool = False,
49) -> dict:
50    """Given a dataclass object, return a json-friendly dict.
51
52    All values will be checked to ensure they match the types specified
53    on fields. Note that a limited set of types and data configurations is
54    supported.
55
56    Values with type Any will be checked to ensure they match types supported
57    directly by json. This does not include types such as tuples which are
58    implicitly translated by Python's json module (as this would break
59    the ability to do a lossless round-trip with data).
60
61    If coerce_to_float is True, integer values present on float typed fields
62    will be converted to float in the dict output. If False, a TypeError
63    will be triggered.
64    """
65
66    out = _Outputter(
67        obj,
68        create=True,
69        codec=codec,
70        coerce_to_float=coerce_to_float,
71        discard_extra_attrs=discard_extra_attrs,
72    ).run()
73    assert isinstance(out, dict)
74    return out

77def dataclass_to_json(
78    obj: Any,
79    coerce_to_float: bool = True,
80    pretty: bool = False,
81    sort_keys: bool | None = None,
82) -> str:
83    """Utility function; return a json string from a dataclass instance.
84
85    Basically json.dumps(dataclass_to_dict(...)).
86    By default, keys are sorted for pretty output and not otherwise, but
87    this can be overridden by supplying a value for the 'sort_keys' arg.
88    """
89
90    jdict = dataclass_to_dict(
91        obj=obj, coerce_to_float=coerce_to_float, codec=Codec.JSON
92    )
93    if sort_keys is None:
94        sort_keys = pretty
95    if pretty:
96        return json.dumps(jdict, indent=2, sort_keys=sort_keys)
97    return json.dumps(jdict, separators=(',', ':'), sort_keys=sort_keys)

166def dataclass_validate(
167    obj: Any,
168    coerce_to_float: bool = True,
169    codec: Codec = Codec.JSON,
170    discard_extra_attrs: bool = False,
171) -> None:
172    """Ensure that values in a dataclass instance are the correct types."""
173
174    # Simply run an output pass but tell it not to generate data;
175    # only run validation.
176    _Outputter(
177        obj,
178        create=False,
179        codec=codec,
180        coerce_to_float=coerce_to_float,
181        discard_extra_attrs=discard_extra_attrs,
182    ).run()

185def dataclass_hash(obj: Any, coerce_to_float: bool = True) -> str:
186    """Calculate a hash for the provided dataclass.
187
188    Basically this emits json for the dataclass (with keys sorted
189    to keep things deterministic) and hashes the resulting string.
190    """
191    import hashlib
192    from base64 import urlsafe_b64encode
193
194    json_dict = dataclass_to_dict(
195        obj, codec=Codec.JSON, coerce_to_float=coerce_to_float
196    )
197
198    # Need to sort keys to keep things deterministic.
199    json_str = json.dumps(json_dict, separators=(',', ':'), sort_keys=True)
200
201    sha = hashlib.sha256()
202    sha.update(json_str.encode())
203
204    # Go with urlsafe base64 instead of the usual hex to save some
205    # space, and kill those ugly padding chars at the end.
206    return urlsafe_b64encode(sha.digest()).decode().strip('=')

47def ioprep(cls: type, globalns: dict | None = None) -> None:
48    """Prep a dataclass type for use with this module's functionality.
49
50    Prepping ensures that all types contained in a data class as well as
51    the usage of said types are supported by this module and pre-builds
52    necessary constructs needed for encoding/decoding/etc.
53
54    Prepping will happen on-the-fly as needed, but a warning will be
55    emitted in such cases, as it is better to explicitly prep all used types
56    early in a process to ensure any invalid types or configuration are caught
57    immediately.
58
59    Prepping a dataclass involves evaluating its type annotations, which,
60    as of PEP 563, are stored simply as strings. This evaluation is done
61    with localns set to the class dict (so that types defined in the class
62    can be used) and globalns set to the containing module's class.
63    It is possible to override globalns for special cases such as when
64    prepping happens as part of an execed string instead of within a
65    module.
66    """
67    PrepSession(explicit=True, globalns=globalns).prep_dataclass(
68        cls, recursion_level=0
69    )

72def ioprepped(cls: type[T]) -> type[T]:
73    """Class decorator for easily prepping a dataclass at definition time.
74
75    Note that in some cases it may not be possible to prep a dataclass
76    immediately (such as when its type annotations refer to forward-declared
77    types). In these cases, dataclass_prep() should be explicitly called for
78    the class as soon as possible; ideally at module import time to expose any
79    errors as early as possible in execution.
80    """
81    ioprep(cls)
82    return cls

102def is_ioprepped_dataclass(obj: Any) -> bool:
103    """Return whether the obj is an ioprepped dataclass type or instance."""
104    cls = obj if isinstance(obj, type) else type(obj)
105    return dataclasses.is_dataclass(cls) and hasattr(cls, PREP_ATTR)

85def will_ioprep(cls: type[T]) -> type[T]:
86    """Class decorator hinting that we will prep a class later.
87
88    In some cases (such as recursive types) we cannot use the @ioprepped
89    decorator and must instead call ioprep() explicitly later. However,
90    some of our custom pylint checking behaves differently when the
91    @ioprepped decorator is present, in that case requiring type annotations
92    to be present and not simply forward declared under an "if TYPE_CHECKING"
93    block. (since they are used at runtime).
94
95    The @will_ioprep decorator triggers the same pylint behavior
96    differences as @ioprepped (which are necessary for the later ioprep() call
97    to work correctly) but without actually running any prep itself.
98    """
99    return cls