feat: make dataclasses automatically de-serializable

Author: ruancomelli

boiling_learning/automl/tuning.py

@@ -6,9 +6,8 @@
 from boiling_learning.datasets.datasets import DatasetTriplet
 from boiling_learning.describe.described import Described
 from boiling_learning.io import json
-from boiling_learning.io.storage import register_deserializer_for_dataclass
+from boiling_learning.io.storage import dataclass
 from boiling_learning.model.model import Evaluation, ModelArchitecture
-from boiling_learning.utils.dataclasses import dataclass
 
 
 @dataclass(frozen=True)
@@ -17,7 +16,6 @@ class TuneModelParams:
     batch_size: int
 
 
-@register_deserializer_for_dataclass
 @dataclass(frozen=True)
 class TuneModelReturn:
     model: ModelArchitecture

boiling_learning/daq/devices.py

@@ -1,6 +1,6 @@
 from nidaqmx.task import Task
 
-from boiling_learning.utils.dataclasses import dataclass
+from boiling_learning.io.storage import dataclass
 
 
 @dataclass

boiling_learning/data/boiling_curve.py

@@ -3,7 +3,7 @@
 from pint import Quantity
 from typing_extensions import Literal
 
-from boiling_learning.utils.dataclasses import dataclass
+from boiling_learning.io.storage import dataclass
 from boiling_learning.utils.units import unit_registry as ureg
 
 Q_ = ureg.Quantity

boiling_learning/datasets/datasets.py

@@ -1,13 +1,10 @@
 from fractions import Fraction
-from pathlib import Path
-from typing import Any, Generic, Optional, TypeVar
+from typing import Generic, Optional, TypeVar
 
 import funcy
 from typing_extensions import NamedTuple
 
-from boiling_learning.io.storage import Metadata, deserialize, load, save, serialize
-from boiling_learning.utils.dataclasses import dataclass
-from boiling_learning.utils.pathutils import resolve
+from boiling_learning.io.storage import dataclass
 
 _T = TypeVar('_T')
 
@@ -57,23 +54,3 @@ def __post_init__(self) -> None:
 
         if not (0 < self.train < 1 and 0 <= self.val < 1 and 0 < self.test < 1):
             raise ValueError('it is required that 0 < (*train*, *test*) < 1 and 0 <= *val* < 1')
-
-
-@serialize.instance(DatasetTriplet)
-def _serialize_dataset_triplet(instance: DatasetTriplet[Any], path: Path) -> None:
-    path = resolve(path, dir=True)
-
-    ds_train, ds_val, ds_test = instance
-
-    save(ds_train, path / 'train')
-    save(ds_val, path / 'val')
-    save(ds_test, path / 'test')
-
-
-@deserialize.dispatch(DatasetTriplet)
-def _deserialize_dataset_triplet(path: Path, metadata: Metadata) -> DatasetTriplet[Any]:
-    ds_train = load(path / 'train')
-    ds_val = load(path / 'val')
-    ds_test = load(path / 'test')
-
-    return DatasetTriplet(ds_train, ds_val, ds_test)

boiling_learning/io/storage.py

@@ -1,9 +1,11 @@
 from __future__ import annotations
 
 import itertools
+from dataclasses import dataclass as _dataclass
 from datetime import timedelta
+from functools import wraps
 from pathlib import Path
-from typing import Any, Type, TypeVar
+from typing import Any, Callable, Type, TypeVar
 
 from classes import AssociatedType, Supports, typeclass
 from typing_extensions import final
@@ -15,6 +17,8 @@
 
 Metadata = json.JSONDataType
 _DataClassType = TypeVar('_DataClassType', bound=Type[DataClass])
+_Any = TypeVar('_Any')
+_AnyCallable = TypeVar('_AnyCallable', bound=Callable[..., Any])
 
 
 @final
@@ -251,3 +255,16 @@ def _deserialize(path: Path, metadata: Metadata) -> DataClass:
         return dataclass_type(**fields)
 
     return dataclass_type
+
+
+def _identity_compose(
+    identity_transform: Callable[[_Any], _Any], function: _AnyCallable
+) -> _AnyCallable:
+    @wraps(function)
+    def _wrapped(*args, **kwargs) -> Any:
+        return identity_transform(function(*args, **kwargs))
+
+    return _wrapped
+
+
+dataclass = _identity_compose(register_deserializer_for_dataclass, _dataclass)

boiling_learning/model/model.py

@@ -100,8 +100,8 @@ def _deserialize_model(path: Path, _metadata: Metadata) -> ModelArchitecture:
 def model_memory_usage_in_bytes(
     architecture: ModelArchitecture, *, batch_size: int
 ) -> Quantity[int]:
-    """
-    Return the estimated memory usage of a given Keras model in bytes.
+    """Return the estimated memory usage of a given Keras model in bytes.
+
     This includes the model weights and layers, but excludes the dataset.
 
     The model shapes are multipled by the batch size, but the weights are not.
@@ -115,7 +115,6 @@ def model_memory_usage_in_bytes(
             pass `1` as the argument here.
     Returns:
         An estimate of the Keras model's memory usage in bytes.
-
     """
     model = architecture.model
 

boiling_learning/model/training.py

@@ -16,10 +16,10 @@
 from boiling_learning.describe.described import Described
 from boiling_learning.describe.describers import describe
 from boiling_learning.io import json
-from boiling_learning.io.storage import Metadata, deserialize, load, save, serialize
+from boiling_learning.io.storage import Metadata, dataclass, deserialize, load, save, serialize
 from boiling_learning.model.callbacks import RegisterEpoch, SaveHistory
 from boiling_learning.model.model import Evaluation, ModelArchitecture
-from boiling_learning.utils.dataclasses import dataclass, fields, shallow_asdict
+from boiling_learning.utils.dataclasses import fields, shallow_asdict
 from boiling_learning.utils.pathutils import resolve
 from boiling_learning.utils.timing import Timer
 from boiling_learning.utils.typeutils import typename
@@ -91,6 +91,8 @@ class FitModelReturn:
     evaluation: Evaluation
 
 
+# NOTE: after v0.34.23, dataclasses are automatically serializable and de-serializable.
+# however, we are keeping this overload for backwards compatibility.
 @serialize.instance(FitModelReturn)
 def _serialize_fit_model_return(instance: FitModelReturn, path: Path) -> None:
     path = resolve(path, dir=True)

boiling_learning/preprocessing/experiment_video.py

@@ -10,8 +10,9 @@
 from boiling_learning.datasets.sliceable import SliceableDataset
 from boiling_learning.describe.describers import describe
 from boiling_learning.io import json
+from boiling_learning.io.storage import dataclass
 from boiling_learning.preprocessing.video import Video, VideoFrame, convert_video
-from boiling_learning.utils.dataclasses import dataclass, field
+from boiling_learning.utils.dataclasses import field
 from boiling_learning.utils.pathutils import PathLike, resolve
 
 

boiling_learning/preprocessing/image_datasets.py

@@ -6,9 +6,10 @@
 
 from boiling_learning.describe.describers import describe
 from boiling_learning.io import json
+from boiling_learning.io.storage import dataclass
 from boiling_learning.preprocessing.experiment_video import ExperimentVideo
 from boiling_learning.utils.collections import KeyedSet
-from boiling_learning.utils.dataclasses import dataclass, dataclass_from_mapping
+from boiling_learning.utils.dataclasses import dataclass_from_mapping
 from boiling_learning.utils.pathutils import PathLike, resolve
 
 

boiling_learning/utils/dataclasses.py

@@ -1,12 +1,11 @@
 import typing
-from dataclasses import asdict, dataclass, field, fields, is_dataclass
+from dataclasses import asdict, field, fields, is_dataclass
 from typing import Any, Callable, Dict, Mapping, Optional, Type, TypeVar, Union
 
 from typing_extensions import TypeGuard
 
 __all__ = (
     'asdict',
-    'dataclass',
     'field',
     'fields',
     'is_dataclass',

main.py

@@ -40,7 +40,7 @@
 )
 from boiling_learning.describe.described import Described
 from boiling_learning.io import json
-from boiling_learning.io.storage import load, save
+from boiling_learning.io.storage import dataclass, load, save
 from boiling_learning.management.allocators import default_table_allocator
 from boiling_learning.management.cacher import CachedFunction, Cacher
 from boiling_learning.model.callbacks import (
@@ -74,7 +74,6 @@
     set_condensation_datasets_data,
 )
 from boiling_learning.scripts.utils.initialization import check_all_paths_exist
-from boiling_learning.utils.dataclasses import dataclass
 from boiling_learning.utils.functional import P
 from boiling_learning.utils.lazy import Lazy, LazyCallable
 from boiling_learning.utils.pathutils import resolve

tests/test_utils/test_dataclasses.py

@@ -1,5 +1,6 @@
+from dataclasses import dataclass
+
 from boiling_learning.utils.dataclasses import (
-    dataclass,
     is_dataclass,
     is_dataclass_class,
     is_dataclass_instance,