Coverage for src/loman/serialization/values.py: 98%
226 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-09-07 00:32 +0000
1"""Transformers for the everyday value types a computation holds.
3These cover the types that reach a node value in practice but that the original
4transformer set did not handle: dates and times, pandas indexes, numpy scalars,
5sets, bytes and decimals. Before this module a DataFrame with a ``DatetimeIndex``
6could not be serialized at all, because the index was encoded element by element
7and no transformer claimed :class:`pandas.Timestamp`.
9Two encoding principles run through the module:
11*Exactness over readability for temporal types.* A timestamp is stored as an
12integer nanosecond count plus a timezone rather than as a formatted string, so a
13round-trip is bit-exact and does not depend on parsing rules.
15*Whole indexes, not element sequences.* An index is encoded as an index, which
16keeps a ``MultiIndex`` a ``MultiIndex`` and turns the default ``RangeIndex`` of a
17100k-row frame into four numbers instead of 100k of them.
18"""
20import base64
21import contextlib
22import datetime
23import decimal
24from collections.abc import Iterable
25from typing import TYPE_CHECKING, Any
27import numpy as np
28import pandas as pd
30from .transformer import CustomTransformer
32if TYPE_CHECKING:
33 from .transformer import Transformer
36class DateTimeTransformer(CustomTransformer):
37 """Transformer for :class:`datetime.datetime`.
39 Encoded as an ISO 8601 string, which carries a UTC offset for aware values.
40 The IANA zone name is kept alongside it when there is one, because an offset
41 alone cannot distinguish ``Europe/London`` in winter from ``UTC``.
42 """
44 @property
45 def name(self) -> str:
46 """Return transformer name."""
47 return "datetime"
49 def to_dict(self, transformer: "Transformer", o: object) -> dict[str, Any]:
50 """Encode a datetime as ISO 8601, keeping the zone name when present."""
51 assert isinstance(o, datetime.datetime) # noqa: S101
52 d: dict[str, Any] = {"iso": o.isoformat()}
53 tzname = getattr(o.tzinfo, "key", None) # zoneinfo.ZoneInfo
54 if tzname is not None:
55 d["tz"] = tzname
56 return d
58 def from_dict(self, transformer: "Transformer", d: dict[str, Any]) -> object:
59 """Reconstruct a datetime, restoring its named zone when recorded."""
60 dt = datetime.datetime.fromisoformat(d["iso"])
61 tzname = d.get("tz")
62 if tzname is not None:
63 from zoneinfo import ZoneInfo
65 dt = dt.astimezone(ZoneInfo(tzname))
66 return dt
68 @property
69 def supported_direct_types(self) -> Iterable[type]:
70 """Return the exact type handled; pandas Timestamps go elsewhere."""
71 return [datetime.datetime]
74class DateTransformer(CustomTransformer):
75 """Transformer for :class:`datetime.date`."""
77 @property
78 def name(self) -> str:
79 """Return transformer name."""
80 return "date"
82 def to_dict(self, transformer: "Transformer", o: object) -> dict[str, Any]:
83 """Encode a date as an ISO 8601 string."""
84 assert isinstance(o, datetime.date) # noqa: S101
85 return {"iso": o.isoformat()}
87 def from_dict(self, transformer: "Transformer", d: dict[str, Any]) -> object:
88 """Reconstruct a date from its ISO 8601 string."""
89 return datetime.date.fromisoformat(d["iso"])
91 @property
92 def supported_direct_types(self) -> Iterable[type]:
93 """Return the exact type handled."""
94 return [datetime.date]
97class TimeTransformer(CustomTransformer):
98 """Transformer for :class:`datetime.time`."""
100 @property
101 def name(self) -> str:
102 """Return transformer name."""
103 return "time"
105 def to_dict(self, transformer: "Transformer", o: object) -> dict[str, Any]:
106 """Encode a time as an ISO 8601 string."""
107 assert isinstance(o, datetime.time) # noqa: S101
108 return {"iso": o.isoformat()}
110 def from_dict(self, transformer: "Transformer", d: dict[str, Any]) -> object:
111 """Reconstruct a time from its ISO 8601 string."""
112 return datetime.time.fromisoformat(d["iso"])
114 @property
115 def supported_direct_types(self) -> Iterable[type]:
116 """Return the exact type handled."""
117 return [datetime.time]
120class TimeDeltaTransformer(CustomTransformer):
121 """Transformer for :class:`datetime.timedelta`.
123 Stored as the three components the type is normalised into, rather than as
124 total seconds, which would lose microsecond precision over long spans.
125 """
127 @property
128 def name(self) -> str:
129 """Return transformer name."""
130 return "timedelta"
132 def to_dict(self, transformer: "Transformer", o: object) -> dict[str, Any]:
133 """Encode a timedelta as its normalised day/second/microsecond parts."""
134 assert isinstance(o, datetime.timedelta) # noqa: S101
135 return {"days": o.days, "seconds": o.seconds, "microseconds": o.microseconds}
137 def from_dict(self, transformer: "Transformer", d: dict[str, Any]) -> object:
138 """Reconstruct a timedelta from its parts."""
139 return datetime.timedelta(days=d["days"], seconds=d["seconds"], microseconds=d["microseconds"])
141 @property
142 def supported_direct_types(self) -> Iterable[type]:
143 """Return the exact type handled."""
144 return [datetime.timedelta]
147class TimestampTransformer(CustomTransformer):
148 """Transformer for :class:`pandas.Timestamp`.
150 :class:`pandas.Timestamp` subclasses :class:`datetime.datetime`, so it needs
151 its own registration: exact-type dispatch would otherwise never reach the
152 datetime transformer, and going through ISO strings would lose nanoseconds.
153 The integer ``value`` is nanoseconds since the epoch, always UTC.
154 """
156 @property
157 def name(self) -> str:
158 """Return transformer name."""
159 return "timestamp"
161 def to_dict(self, transformer: "Transformer", o: object) -> dict[str, Any]:
162 """Encode a Timestamp as epoch nanoseconds plus timezone and resolution.
164 ``Timestamp.value`` is always nanoseconds, but the resolution the value
165 was held at is a separate property and is preserved so that a
166 microsecond timestamp does not come back claiming nanosecond precision.
167 """
168 assert isinstance(o, pd.Timestamp) # noqa: S101
169 d: dict[str, Any] = {"value": o.value, "unit": o.unit}
170 if o.tz is not None:
171 d["tz"] = str(o.tz)
172 return d
174 def from_dict(self, transformer: "Transformer", d: dict[str, Any]) -> object:
175 """Reconstruct a Timestamp, re-applying its timezone and resolution."""
176 ts = pd.Timestamp(d["value"])
177 tz = d.get("tz")
178 if tz is not None:
179 ts = ts.tz_localize("UTC").tz_convert(tz)
180 unit = d.get("unit")
181 if unit is not None:
182 ts = ts.as_unit(unit)
183 return ts
185 @property
186 def supported_direct_types(self) -> Iterable[type]:
187 """Return the exact type handled."""
188 return [pd.Timestamp]
191class PandasTimedeltaTransformer(CustomTransformer):
192 """Transformer for :class:`pandas.Timedelta`, stored as nanoseconds."""
194 @property
195 def name(self) -> str:
196 """Return transformer name."""
197 return "pd_timedelta"
199 def to_dict(self, transformer: "Transformer", o: object) -> dict[str, Any]:
200 """Encode a pandas Timedelta as a nanosecond count plus its resolution."""
201 assert isinstance(o, pd.Timedelta) # noqa: S101
202 return {"value": o.value, "unit": o.unit}
204 def from_dict(self, transformer: "Transformer", d: dict[str, Any]) -> object:
205 """Reconstruct a pandas Timedelta from its nanosecond count."""
206 td = pd.Timedelta(d["value"])
207 unit = d.get("unit")
208 if unit is not None:
209 td = td.as_unit(unit)
210 return td
212 @property
213 def supported_direct_types(self) -> Iterable[type]:
214 """Return the exact type handled."""
215 return [pd.Timedelta]
218class NaTTransformer(CustomTransformer):
219 """Transformer for :data:`pandas.NaT`.
221 ``NaT`` is a singleton of its own type rather than a Timestamp, so nothing
222 else claims it, and it would otherwise fail as an unknown type inside any
223 datetime column holding a missing value.
224 """
226 @property
227 def name(self) -> str:
228 """Return transformer name."""
229 return "nat"
231 def to_dict(self, transformer: "Transformer", o: object) -> dict[str, Any]:
232 """Encode NaT as an empty marker."""
233 return {}
235 def from_dict(self, transformer: "Transformer", d: dict[str, Any]) -> object:
236 """Return the NaT singleton."""
237 return pd.NaT
239 @property
240 def supported_direct_types(self) -> Iterable[type]:
241 """Return the NaT singleton's type."""
242 return [type(pd.NaT)]
245class NumpyScalarTransformer(CustomTransformer):
246 """Transformer for numpy scalar types (:class:`numpy.generic`).
248 Registered against the ``np.generic`` base so every width is covered by one
249 transformer. The dtype string is kept so ``np.int32`` does not come back as
250 ``np.int64``, and ``np.float64`` does not silently degrade to a plain float.
251 """
253 @property
254 def name(self) -> str:
255 """Return transformer name."""
256 return "npscalar"
258 def to_dict(self, transformer: "Transformer", o: object) -> dict[str, Any]:
259 """Encode a numpy scalar as its dtype plus a plain Python value."""
260 assert isinstance(o, np.generic) # noqa: S101
261 dtype = o.dtype
262 if dtype.kind in "Mm":
263 # datetime64 / timedelta64: the integer tick count is exact, and
264 # .item() would hand back a datetime whose unit had been forgotten.
265 return {"dtype": dtype.str, "value": int(o.view("int64"))}
266 return {"dtype": dtype.str, "value": transformer.to_dict(o.item())}
268 def from_dict(self, transformer: "Transformer", d: dict[str, Any]) -> object:
269 """Reconstruct a numpy scalar of the recorded dtype."""
270 dtype = np.dtype(d["dtype"])
271 if dtype.kind in "Mm":
272 return np.int64(d["value"]).view(dtype)
273 return dtype.type(transformer.from_dict(d["value"]))
275 @property
276 def supported_subtypes(self) -> Iterable[Any]:
277 """Match every numpy scalar type."""
278 return [np.generic]
281class SetTransformer(CustomTransformer):
282 """Transformer for :class:`set` and :class:`frozenset`.
284 Members are sorted by their encoded form where possible so that two equal
285 sets serialize identically --- without it, a saved file would differ run to
286 run with hash randomisation, defeating byte-level comparison of two saves.
287 """
289 @property
290 def name(self) -> str:
291 """Return transformer name."""
292 return "set"
294 def to_dict(self, transformer: "Transformer", o: object) -> dict[str, Any]:
295 """Encode a set as a list of encoded members plus its mutability."""
296 assert isinstance(o, (set, frozenset)) # noqa: S101
297 values = [transformer.to_dict(x) for x in o]
298 with contextlib.suppress(TypeError): # pragma: no cover - unorderable encodings
299 values.sort(key=_sort_key)
300 return {"values": values, "frozen": isinstance(o, frozenset)}
302 def from_dict(self, transformer: "Transformer", d: dict[str, Any]) -> object:
303 """Reconstruct a set or frozenset from its members."""
304 values = (transformer.from_dict(x) for x in d["values"])
305 return frozenset(values) if d.get("frozen") else set(values)
307 @property
308 def supported_direct_types(self) -> Iterable[type]:
309 """Return the exact types handled."""
310 return [set, frozenset]
313def _sort_key(encoded: Any) -> tuple[str, str]:
314 """Return a total ordering key for an encoded value.
316 Encoded members can be scalars or dicts, which do not compare with one
317 another, so ordering falls back to the type name and the repr.
318 """
319 return (type(encoded).__name__, repr(encoded))
322class BytesTransformer(CustomTransformer):
323 """Transformer for :class:`bytes` and :class:`bytearray`, as base64."""
325 @property
326 def name(self) -> str:
327 """Return transformer name."""
328 return "bytes"
330 def to_dict(self, transformer: "Transformer", o: object) -> dict[str, Any]:
331 """Encode a byte string as base64 ASCII."""
332 assert isinstance(o, (bytes, bytearray)) # noqa: S101
333 return {"b64": base64.b64encode(bytes(o)).decode("ascii"), "mutable": isinstance(o, bytearray)}
335 def from_dict(self, transformer: "Transformer", d: dict[str, Any]) -> object:
336 """Reconstruct a byte string from base64."""
337 raw = base64.b64decode(d["b64"].encode("ascii"))
338 return bytearray(raw) if d.get("mutable") else raw
340 @property
341 def supported_direct_types(self) -> Iterable[type]:
342 """Return the exact types handled."""
343 return [bytes, bytearray]
346class DecimalTransformer(CustomTransformer):
347 """Transformer for :class:`decimal.Decimal`, stored as its exact string."""
349 @property
350 def name(self) -> str:
351 """Return transformer name."""
352 return "decimal"
354 def to_dict(self, transformer: "Transformer", o: object) -> dict[str, Any]:
355 """Encode a Decimal as the string that reproduces it exactly."""
356 assert isinstance(o, decimal.Decimal) # noqa: S101
357 return {"value": str(o)}
359 def from_dict(self, transformer: "Transformer", d: dict[str, Any]) -> object:
360 """Reconstruct a Decimal from its string form."""
361 return decimal.Decimal(d["value"])
363 @property
364 def supported_direct_types(self) -> Iterable[type]:
365 """Return the exact type handled."""
366 return [decimal.Decimal]
369class IndexTransformer(CustomTransformer):
370 """Transformer for :class:`pandas.Index` and its subclasses.
372 Encoding an index as an index rather than as a list of its elements is what
373 keeps a ``MultiIndex`` a ``MultiIndex`` --- previously it came back as a flat
374 ``Index`` of tuples --- and what stops the default ``RangeIndex`` of a large
375 frame from being written out one integer at a time.
377 Five shapes are recognised, discriminated by ``kind``: ``range``,
378 ``datetime``, ``timedelta``, ``multi`` and ``base``. A ``MultiIndex``
379 encodes its levels recursively as indexes, so a level that is itself a
380 ``DatetimeIndex`` keeps its type and timezone.
381 """
383 @property
384 def name(self) -> str:
385 """Return transformer name."""
386 return "index"
388 def to_dict(self, transformer: "Transformer", o: object) -> dict[str, Any]:
389 """Encode an index according to its concrete pandas type."""
390 assert isinstance(o, pd.Index) # noqa: S101
391 if isinstance(o, pd.MultiIndex):
392 return {
393 "kind": "multi",
394 "levels": [transformer.to_dict(level) for level in o.levels],
395 "codes": [list(map(int, codes)) for codes in o.codes],
396 "names": list(o.names),
397 }
398 if isinstance(o, pd.RangeIndex):
399 return {
400 "kind": "range",
401 "start": int(o.start),
402 "stop": int(o.stop),
403 "step": int(o.step),
404 "name": transformer.to_dict(o.name),
405 }
406 if isinstance(o, pd.DatetimeIndex):
407 # asi8 counts ticks in the index's own resolution, which is
408 # microseconds by default in pandas 3 and nanoseconds before it.
409 # Recording the unit is what keeps the values from being reread at
410 # the wrong scale.
411 return {
412 "kind": "datetime",
413 "values": transformer.to_dict(o.asi8),
414 "unit": o.unit,
415 "tz": str(o.tz) if o.tz is not None else None,
416 "freq": o.freqstr if o.freq is not None else None,
417 "name": transformer.to_dict(o.name),
418 }
419 if isinstance(o, pd.TimedeltaIndex):
420 return {
421 "kind": "timedelta",
422 "values": transformer.to_dict(o.asi8),
423 "unit": o.unit,
424 "freq": o.freqstr if o.freq is not None else None,
425 "name": transformer.to_dict(o.name),
426 }
427 return {
428 "kind": "base",
429 "dtype": str(o.dtype),
430 "values": _encode_index_values(transformer, o),
431 "name": transformer.to_dict(o.name),
432 }
434 def from_dict(self, transformer: "Transformer", d: dict[str, Any]) -> object:
435 """Reconstruct an index of the recorded kind."""
436 kind = d["kind"]
437 if kind == "multi":
438 levels = [transformer.from_dict(level) for level in d["levels"]]
439 return pd.MultiIndex(levels=levels, codes=d["codes"], names=d["names"])
440 name = transformer.from_dict(d.get("name"))
441 if kind == "range":
442 return pd.RangeIndex(start=d["start"], stop=d["stop"], step=d["step"], name=name)
443 if kind == "datetime":
444 unit = d.get("unit", "ns")
445 values = np.asarray(transformer.from_dict(d["values"]), dtype="int64").view(f"datetime64[{unit}]")
446 idx = pd.DatetimeIndex(values, name=name)
447 tz = d.get("tz")
448 if tz is not None:
449 idx = idx.tz_localize("UTC").tz_convert(tz)
450 return _restore_freq(idx, d.get("freq"))
451 if kind == "timedelta":
452 unit = d.get("unit", "ns")
453 values = np.asarray(transformer.from_dict(d["values"]), dtype="int64").view(f"timedelta64[{unit}]")
454 return _restore_freq(pd.TimedeltaIndex(values, name=name), d.get("freq"))
455 return pd.Index(transformer.from_dict(d["values"]), dtype=d["dtype"], name=name)
457 @property
458 def supported_subtypes(self) -> Iterable[Any]:
459 """Match every pandas Index subclass."""
460 return [pd.Index]
463def _encode_index_values(transformer: "Transformer", index: "pd.Index") -> Any:
464 """Encode a plain index's values, as an array where that is lossless.
466 Same reasoning as for DataFrame columns: a numpy-backed index goes through
467 the ndarray transformer, which lets a large one be written out of line
468 rather than one number at a time in the manifest. Extension dtypes go
469 element-wise, since their numpy form is not the same thing.
470 """
471 dtype = index.dtype
472 if isinstance(dtype, np.dtype) and not dtype.hasobject:
473 return transformer.to_dict(index.to_numpy())
474 return transformer.to_dict(index.tolist())
477def _restore_freq(index: Any, freq: str | None) -> Any:
478 """Re-apply a recorded frequency to *index*, ignoring one that no longer fits.
480 ``freq`` is a derived property: it describes a regular spacing the values
481 already have. Setting it cannot change the data, and a value that does not
482 match --- which pandas rejects --- is better dropped than fatal.
483 """
484 if freq is None:
485 return index
486 with contextlib.suppress(ValueError, TypeError): # pragma: no cover - defensive
487 index.freq = freq
488 return index
491VALUE_TRANSFORMERS: list[type[CustomTransformer]] = [
492 DateTimeTransformer,
493 DateTransformer,
494 TimeTransformer,
495 TimeDeltaTransformer,
496 TimestampTransformer,
497 PandasTimedeltaTransformer,
498 NaTTransformer,
499 NumpyScalarTransformer,
500 SetTransformer,
501 BytesTransformer,
502 DecimalTransformer,
503 IndexTransformer,
504]
507def register_value_transformers(t: "Transformer") -> None:
508 """Register every everyday-value transformer on *t*."""
509 for transformer_cls in VALUE_TRANSFORMERS:
510 t.register(transformer_cls())