Add to_dict and from_dict methods to Labware

Author: michaelosthege

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "robotools"
-version = "1.14.0"
+version = "1.15.0"
 description = "Pythonic in-silico liquid handling and creation of Tecan FreedomEVO worklists."
 readme = "README.md"
 requires-python = ">=3.10"

robotools/liquidhandling/labware.py

@@ -3,7 +3,7 @@
 
 import logging
 import warnings
-from typing import Dict, List, Literal, Mapping, Optional, Sequence, Tuple, Union
+from typing import Any, Dict, List, Literal, Mapping, Optional, Sequence, Tuple, Union
 
 import numpy as np
 
@@ -369,6 +369,104 @@ def __repr__(self) -> str:
     def __str__(self) -> str:
         return self.__repr__()
 
+    def to_dict(self) -> dict[str, Any]:
+        """Serialize to a dict of only atomic data types.
+
+        Use :meth:`Labware.from_dict()` to recreate a labware from the dict produced by this method."""
+        return {
+            "name": self.name,
+            "rows": self.n_rows,
+            "columns": self.n_columns,
+            "min_volume": self.min_volume,
+            "max_volume": self.max_volume,
+            "volumes": self.volumes.tolist(),
+            "history": [h.tolist() for h in self._history],
+            "labels": self._labels,
+            "composition": {k: v.tolist() for k, v in self.composition.items()},
+            "is_trough": self.is_trough,
+            "virtual_rows": self.virtual_rows,
+        }
+
+    @classmethod
+    def _from_dict_init_kwargs(cls, data: Mapping[str, Any]) -> Dict[str, Any]:
+        """Extract initializer parameters from a dict created by :meth:`Labware.to_dict()`."""
+        return {
+            "name": data["name"],
+            "rows": data["rows"],
+            "columns": data["columns"],
+            "min_volume": data["min_volume"],
+            "max_volume": data["max_volume"],
+        }
+
+    @classmethod
+    def _from_dict_attributes(cls, data: Mapping[str, Any], vshape: tuple[int, int]) -> Dict[str, Any]:
+        """Extracts state attributes from a dict created by :meth:`Labware.to_dict()`.
+
+        Parameters
+        ----------
+        data
+            Dict created by :meth:`Labware.to_dict()`.
+        vshape
+            Shape of the volumes array.
+
+        Returns
+        -------
+        Dict of attribute names and values to be assigned to the labware after initialization.
+        """
+
+        def as_2d(arr):
+            return np.asarray(arr, dtype=float).reshape(vshape)
+
+        return {
+            "_volumes": as_2d(data["volumes"]),
+            "_history": [as_2d(h) for h in data["history"]],
+            "_labels": data["labels"],
+            "_composition": {k: as_2d(v) for k, v in data["composition"].items()},
+        }
+
+    @classmethod
+    def from_dict(cls, data: Mapping[str, Any], cls_trough: type | None = None) -> "Labware | Trough":
+        """Creates a labware from a dict produced by :meth:`Labware.to_dict()`.
+
+        Parameters
+        ----------
+        data
+            Dict created by :meth:`Labware.to_dict()`.
+        cls_trough
+            Optional trough type to use in case the data dict corresponds to a trough labware.
+            Defaults to :class:`robotools.Trough` with a warning if the main class is not a `Trough` itself.
+        """
+        # First we must determine the type to be instantiated.
+        # If the data corresponds to a trough, we want to create a trough object, which is OK with the return type hint,
+        # because Trough is a subclass of Labware. Therefore Labware.from_dict can be used to restore Troughs as well.
+        # We allow the caller to specify a derived trough class, or default to the Trough class in this module.
+        # This is useful when the actual type used by the caller is a custom subclass of Labware/Trough.
+        _cls: type[Labware] | type[Trough]
+        if data.get("is_trough", False) and not issubclass(cls, Trough):
+            if cls_trough is None:
+                if cls is Labware:
+                    warnings.warn(
+                        "Data corresponds to a trough, but no trough type was given."
+                        " Defaulting to `robotools.Trough`."
+                        f" Specify a trough type to avoid downcasting {cls} to {Trough}.",
+                        UserWarning,
+                        stacklevel=2,
+                    )
+                _cls = Trough
+            else:
+                _cls = cls_trough
+        else:
+            _cls = cls
+
+        # Now that we have the type we can proceed with extraction of initializer parameters and state attributes from the dict.
+        init_kwargs = _cls._from_dict_init_kwargs(data)
+        lware = _cls(**init_kwargs)
+        attributes = _cls._from_dict_attributes(data, lware.volumes.shape)
+        # assign state attributes
+        for aname, avalue in attributes.items():
+            setattr(lware, aname, avalue)
+        return lware
+
 
 class Trough(Labware):
     """Special type of labware that can be accessed by many pipette tips in parallel."""
@@ -428,3 +526,14 @@ def __init__(
             virtual_rows=virtual_rows,
             component_names=component_names,
         )
+
+    @classmethod
+    def _from_dict_init_kwargs(cls, data: Mapping[str, Any]) -> Dict[str, Any]:
+        """Extract initializer parameters from a dict created by :meth:`Labware.to_dict()`.
+
+        Overrides the superclass method to account for the different initializer signature of troughs.
+        """
+        lware_kwargs = Labware._from_dict_init_kwargs(data)
+        lware_kwargs.pop("rows")
+        lware_kwargs["virtual_rows"] = data["virtual_rows"]
+        return lware_kwargs

robotools/liquidhandling/test_serialization.py

@@ -0,0 +1,153 @@
+import warnings
+
+import numpy as np
+import pytest
+
+from robotools import Labware, Trough
+
+
+def assert_labware_equal(original: Labware, restored: Labware):
+    assert isinstance(restored, Labware)
+    assert len(restored.history) == len(original.history)
+    # compare history entries
+    for (olabel, ovols), (rlabel, rvols) in zip(original.history, restored.history, strict=True):
+        assert rlabel == olabel
+        np.testing.assert_array_equal(ovols, rvols)
+    # compare composition of all wells
+    for w in original.wells.flatten():
+        assert original.get_well_composition(w) == restored.get_well_composition(w)
+    return
+
+
+def test_labware_to_dict_from_dict():
+    # create a labware and change it's state
+    mtp = Labware(
+        "emtepe",
+        2,
+        3,
+        min_volume=30,
+        max_volume=240,
+        initial_volumes=50,
+    )
+    mtp.add("B02", 50, label="add toxic stuff", compositions=[{"toxin": 0.5, "water": 0.5}])
+
+    # encode and recreate
+    mtp_dict = mtp.to_dict()
+    # the dict must not contain numpy arrays
+    # that's why we can equals compare it here.
+    assert mtp_dict == {
+        "name": "emtepe",
+        "rows": 2,
+        "columns": 3,
+        "min_volume": 30,
+        "max_volume": 240,
+        "volumes": [[50.0, 50.0, 50.0], [50.0, 100.0, 50.0]],
+        "history": [[[50.0, 50.0, 50.0], [50.0, 50.0, 50.0]], [[50.0, 50.0, 50.0], [50.0, 100.0, 50.0]]],
+        "labels": ["initial", "add toxic stuff"],
+        "composition": {
+            "emtepe.A01": [[1.0, 0.0, 0.0], [0.0, 0.0, 0.0]],
+            "emtepe.A02": [[0.0, 1.0, 0.0], [0.0, 0.0, 0.0]],
+            "emtepe.A03": [[0.0, 0.0, 1.0], [0.0, 0.0, 0.0]],
+            "emtepe.B01": [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]],
+            "emtepe.B02": [[0.0, 0.0, 0.0], [0.0, 0.5, 0.0]],
+            "emtepe.B03": [[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]],
+            "toxin": [[0.0, 0.0, 0.0], [0.0, 0.25, 0.0]],
+            "water": [[0.0, 0.0, 0.0], [0.0, 0.25, 0.0]],
+        },
+        "is_trough": False,
+        "virtual_rows": None,
+    }
+
+    restored = Labware.from_dict(mtp_dict)
+
+    assert_labware_equal(mtp, restored)
+    assert isinstance(mtp, Labware)
+    pass
+
+
+def test_trough_to_dict_from_dict():
+    # create a labware and change it's state
+    orig = Trough(
+        "buffers",
+        8,
+        2,
+        min_volume=30_000,
+        max_volume=200_000,
+        initial_volumes=[100_000, 50_000],
+        column_names=["left", "right"],
+    )
+    orig.add("A01", 50, label="add toxic stuff", compositions=[{"toxin": 0.5, "water": 0.5}])
+
+    # encode and recreate
+    assert orig.is_trough
+    data = orig.to_dict()
+    assert data["is_trough"] is True
+    restored = Trough.from_dict(data)
+    assert isinstance(restored, Trough)
+
+    assert_labware_equal(orig, restored)
+    pass
+
+
+def test_labware_from_dict_can_return_troughs():
+    # create a labware and change it's state
+    orig = Trough(
+        "buffers",
+        8,
+        2,
+        min_volume=30_000,
+        max_volume=200_000,
+        initial_volumes=[100_000, 50_000],
+        column_names=["left", "right"],
+    )
+    orig.add("A01", 50, label="add toxic stuff", compositions=[{"toxin": 0.5, "water": 0.5}])
+
+    # encode and recreate
+    assert orig.is_trough
+    data = orig.to_dict()
+    assert data["is_trough"] is True
+    with pytest.warns(UserWarning, match="downcasting"):
+        restored = Labware.from_dict(data)
+    assert isinstance(restored, Trough)
+
+    assert_labware_equal(orig, restored)
+    pass
+
+
+def test_from_dict_warns_about_downcasting_trough():
+    class MyTrough(Trough):
+        pass
+
+    # create a labware and change it's state
+    orig = MyTrough(
+        "buffers",
+        8,
+        2,
+        min_volume=30_000,
+        max_volume=200_000,
+        initial_volumes=[100_000, 50_000],
+        column_names=["left", "right"],
+    )
+    orig.add("A01", 50, label="add toxic stuff", compositions=[{"toxin": 0.5, "water": 0.5}])
+
+    # encode and recreate
+    assert orig.is_trough
+    data = orig.to_dict()
+    assert data["is_trough"] is True
+
+    # robotools doesn't know about MyTrough, so we'll get a Trough with a warning
+    with pytest.warns(UserWarning, match="downcasting"):
+        restored = Labware.from_dict(data)
+    # instantiation worked, but the type was downcasted
+    assert isinstance(restored, Trough)
+    # we can avoid downcasting by passing the trough type explicitly
+    with warnings.catch_warnings():
+        warnings.simplefilter("error")
+        restored2 = Labware.from_dict(data, cls_trough=MyTrough)
+    assert isinstance(restored2, MyTrough)
+    # or by calling through the custom type
+    with warnings.catch_warnings():
+        warnings.simplefilter("error")
+        restored3 = MyTrough.from_dict(data)
+    assert isinstance(restored3, MyTrough)
+    pass