adding Pickle mixin (#64)

* adding Pickle mixin
* isinstance
* compress
* object

Author: Borda

src/litmodels/integrations/mixins.py

@@ -0,0 +1,85 @@
+import pickle
+import tempfile
+from abc import ABC
+from pathlib import Path
+from typing import Optional
+
+from litmodels import download_model, upload_model
+
+
+class ModelRegistryMixin(ABC):
+    """Mixin for model registry integration."""
+
+    def push_to_registry(
+        self, model_name: Optional[str] = None, model_version: Optional[str] = None, temp_folder: Optional[str] = None
+    ) -> None:
+        """Push the model to the registry.
+
+        Args:
+            model_name: The name of the model. If not use the class name.
+            model_version: The version of the model. If None, the latest version is used.
+            temp_folder: The temporary folder to save the model. If None, a default temporary folder is used.
+        """
+
+    @classmethod
+    def pull_from_registry(
+        cls, model_name: str, model_version: Optional[str] = None, temp_folder: Optional[str] = None
+    ) -> object:
+        """Pull the model from the registry.
+
+        Args:
+            model_name: The name of the model.
+            model_version: The version of the model. If None, the latest version is used.
+            temp_folder: The temporary folder to save the model. If None, a default temporary folder is used.
+        """
+
+
+class PickleRegistryMixin(ABC):
+    """Mixin for pickle registry integration."""
+
+    def push_to_registry(
+        self, model_name: Optional[str] = None, model_version: Optional[str] = None, temp_folder: Optional[str] = None
+    ) -> None:
+        """Push the model to the registry.
+
+        Args:
+            model_name: The name of the model. If not use the class name.
+            model_version: The version of the model. If None, the latest version is used.
+            temp_folder: The temporary folder to save the model. If None, a default temporary folder is used.
+        """
+        if model_name is None:
+            model_name = self.__class__.__name__
+        if temp_folder is None:
+            temp_folder = tempfile.gettempdir()
+        pickle_path = Path(temp_folder) / f"{model_name}.pkl"
+        with open(pickle_path, "wb") as fp:
+            pickle.dump(self, fp, protocol=pickle.HIGHEST_PROTOCOL)
+        model_registry = f"{model_name}:{model_version}" if model_version else model_name
+        upload_model(name=model_registry, model=pickle_path)
+
+    @classmethod
+    def pull_from_registry(
+        cls, model_name: str, model_version: Optional[str] = None, temp_folder: Optional[str] = None
+    ) -> object:
+        """Pull the model from the registry.
+
+        Args:
+            model_name: The name of the model.
+            model_version: The version of the model. If None, the latest version is used.
+            temp_folder: The temporary folder to save the model. If None, a default temporary folder is used.
+        """
+        if temp_folder is None:
+            temp_folder = tempfile.gettempdir()
+        model_registry = f"{model_name}:{model_version}" if model_version else model_name
+        files = download_model(name=model_registry, download_dir=temp_folder)
+        pkl_files = [f for f in files if f.endswith(".pkl")]
+        if not pkl_files:
+            raise RuntimeError(f"No pickle file found for model: {model_registry} with {files}")
+        if len(pkl_files) > 1:
+            raise RuntimeError(f"Multiple pickle files found for model: {model_registry} with {pkl_files}")
+        pkl_path = Path(temp_folder) / pkl_files[0]
+        with open(pkl_path, "rb") as fp:
+            obj = pickle.load(fp)
+        if not isinstance(obj, cls):
+            raise RuntimeError(f"Unpickled object is not of type {cls.__name__}: {type(obj)}")
+        return obj

tests/integrations/test_mixins.py

@@ -0,0 +1,31 @@
+from unittest import mock
+
+from litmodels.integrations.mixins import PickleRegistryMixin
+
+
+class DummyModel(PickleRegistryMixin):
+    def __init__(self, value):
+        self.value = value
+
+    def __eq__(self, other):
+        return isinstance(other, DummyModel) and self.value == other.value
+
+
+@mock.patch("litmodels.integrations.mixins.upload_model")
+@mock.patch("litmodels.integrations.mixins.download_model")
+def test_pickle_push_and_pull(mock_download_model, mock_upload_model, tmp_path):
+    # Create an instance of DummyModel and call push_to_registry.
+    dummy = DummyModel(42)
+    dummy.push_to_registry(model_version="v1", temp_folder=str(tmp_path))
+    # The expected registry name is "dummy_model:v1" and the file should be placed in the temp folder.
+    expected_path = tmp_path / "DummyModel.pkl"
+    mock_upload_model.assert_called_once_with(name="DummyModel:v1", model=expected_path)
+
+    # Set the mock to return the full path to the pickle file.
+    mock_download_model.return_value = ["DummyModel.pkl"]
+    # Call pull_from_registry and load the DummyModel instance.
+    loaded_dummy = DummyModel.pull_from_registry(
+        model_name="dummy_model", model_version="v1", temp_folder=str(tmp_path)
+    )
+    # Verify that the unpickled instance has the expected value.
+    assert loaded_dummy.value == 42