Release 0.2.0

pyproject.toml

@@ -31,6 +31,7 @@ flake8-docstrings = { version = "^1.6.0", optional = true }
 mypy = {version = "^0.900", optional = true}
 pytest  = { version = "^6.2.4", optional = true}
 pytest-cov  = { version = "^2.12.0", optional = true}
+pytest-mock = { version = "^3.10.0", optional = true}
 tox  = { version = "^3.20.1", optional = true}
 virtualenv  = { version = "^20.2.2", optional = true}
 pip  = { version = "^20.3.1", optional = true}
@@ -54,7 +55,8 @@ test = [
     "mypy",
     "flake8",
     "flake8-docstrings",
-    "pytest-cov"
+    "pytest-cov",
+    "pytest-mock",
     ]
 
 dev = ["tox", "pre-commit", "virtualenv", "pip", "twine", "toml", "bump2version"]

tests/test_report.py

@@ -1,15 +1,26 @@
 #!/usr/bin/env python
 """Tests for `trackable` package."""
 
-from unittest.mock import Mock
+import os
+import tempfile
 
 import numpy as np
 import pandas as pd
 import pytest
 from pandas.io.formats.style import Styler
 
 from trackable import Report
-from trackable.exceptions import ModelAlreadyExistsError
+from trackable.exceptions import ArchiveAlreadyExistsError, ModelAlreadyExistsError, ModelDoesNotExistError
+
+MOCK_MODEL_NAME = "Mock"
+
+
+class MockModel:
+    """Simple model for testing"""
+
+    def predict(self, x):
+        """Mock predict f(x) = x"""
+        return x
 
 
 @pytest.fixture
@@ -40,18 +51,18 @@ def mock_metric_2(x, y):
 
 
 @pytest.fixture
-def mock_model():
+def mock_model(mocker):
     """Fixture to mock a dummy model."""
-    model = Mock()
-    model.predict = lambda x: x
+    model = MockModel()
+    model.__class__.__name__ = MOCK_MODEL_NAME
     return model
 
 
 def test_add_model_1(mock_report, mock_model):
     """Test add_model 1: add a single model."""
     mock_report.add_model(mock_model)
     print(mock_report._results, mock_report._models)
-    assert mock_report._results == [{"<lambda>": 1, "name": "Mock"}]
+    assert mock_report._results == [{"<lambda>": 1, "name": MOCK_MODEL_NAME}]
     assert mock_report._models == {"Mock": mock_model}
 
 
@@ -68,7 +79,7 @@ def test_add_model_3(mock_report_complex, mock_model):
     """Test add model 3: add model with multiple metrics."""
     mock_report_complex.add_model(mock_model)
     print(mock_report_complex._results, mock_report_complex._models)
-    assert mock_report_complex._results == [{"mock_metric_1": 1, "mock_metric_2": 2, "name": "Mock"}]
+    assert mock_report_complex._results == [{"mock_metric_1": 1, "mock_metric_2": 2, "name": MOCK_MODEL_NAME}]
     assert mock_report_complex._models == {"Mock": mock_model}
 
 
@@ -84,7 +95,7 @@ def test_generate_1(mock_report, mock_model):
     """Test generate 1: show correct results in dataframe."""
     mock_report.add_model(mock_model)
     report = mock_report.generate(False)
-    correct = pd.DataFrame([{"<lambda>": 1.0, "name": "Mock"}]).set_index("name")
+    correct = pd.DataFrame([{"<lambda>": 1.0, "name": MOCK_MODEL_NAME}]).set_index("name")
     print(report)
     print(correct)
     assert correct.equals(report)
@@ -121,3 +132,61 @@ def test_generate_4(mock_report, mock_model):
     assert isinstance(report, pd.DataFrame)
     with pytest.raises(TypeError):
         mock_report.generate(highlight="42")
+
+
+def test_get_model_1(mock_report, mock_model):
+    """Test get_model 1: get correct model"""
+    mock_report.add_model(mock_model, "Model 1")
+    model = mock_report.get_model("Model 1")
+    assert mock_model == model
+
+
+def test_get_model_2(mock_report, mock_model):
+    """Test get_model 2: raise an error for incorrect model"""
+    mock_report.add_model(mock_model, "Model 1")
+    with pytest.raises(ModelDoesNotExistError):
+        mock_report.get_model("Model 2")
+
+
+def test_remove_model_1(mock_report, mock_model):
+    """Test remove_model 1: remove correct model"""
+    mock_report.add_model(mock_model, "Model 1")
+    model = mock_report.remove_model("Model 1")
+    assert mock_model == model
+    assert model not in mock_report._models
+
+
+def test_remove_model_2(mock_report, mock_model):
+    """Test remove_model 2: raise an error for non-existant model"""
+    mock_report.add_model(mock_model, "Model 1")
+    with pytest.raises(ModelDoesNotExistError):
+        mock_report.remove_model("Model 2")
+
+
+def test_archive_model_1(mock_report):
+    """Test archive model 1: archive model correctly (0 models)"""
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        mock_report.save(os.path.join(tmp_dir, "test"))
+        assert os.path.isfile(os.path.join(tmp_dir, "test.zip"))
+
+
+def test_archive_model_2(mock_report):
+    """Test archive model 2: raise error for existing archive"""
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        mock_report.save(os.path.join(tmp_dir, "test"))
+        with pytest.raises(ArchiveAlreadyExistsError):
+            mock_report.save(os.path.join(tmp_dir, "test"))
+
+
+def test_load_model_1(mock_report, mock_model):
+    """Test load model 1: save and load a model correctly"""
+    mock_report.add_model(mock_model)
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        test_zip = "test"
+        mock_report.save(os.path.join(tmp_dir, test_zip))
+        X_test = np.array([])
+        y_test = np.array([])
+        metrics = [lambda x, y: 1.0]
+        new_report = Report.load(X_test, y_test, metrics, os.path.join(tmp_dir, test_zip))
+    assert type(new_report._models[MOCK_MODEL_NAME]) is type(mock_model)
+    assert new_report._results == mock_report._results

trackable/exceptions.py

@@ -3,3 +3,15 @@
 
 class ModelAlreadyExistsError(Exception):
     """Raise error if the model already exists. This is to avoid duplicates."""
+
+
+class ModelDoesNotExistError(KeyError):
+    """Raises an error if the model does not exist. Used when fetching models."""
+
+
+class ArchiveAlreadyExistsError(FileExistsError):
+    """Raise an error if the report archive already exists. This is to avoid overwriting"""
+
+
+class ArchiveDoesNotExistError(FileNotFoundError):
+    """Raises an error if the report archive does not exist."""

trackable/report.py

@@ -1,15 +1,21 @@
 """Module used to generate a minimal report to track ML models."""
 
-from typing import Any, Dict, List, Literal, Optional, Union
+import os
+import pickle
+import shutil
+import tempfile
+from typing import Any, Dict, List, Literal, Optional, Type, TypeVar, Union
 
 import pandas as pd
 from pandas.io.formats.style import Styler
 
 from trackable import types
-from trackable.exceptions import ModelAlreadyExistsError
+from trackable.exceptions import ArchiveAlreadyExistsError, ModelAlreadyExistsError, ModelDoesNotExistError
 
 __all__ = ["Report"]
 
+R = TypeVar("R", bound="Report")
+
 
 class Report:
     """A minimalistic model reporting class.
@@ -44,6 +50,50 @@ def __init__(self, X_test: Any, y_test: Any, metrics: List[types.Metric]) -> Non
         self._models: dict = {}
         self._results: List[dict] = []
 
+    @classmethod
+    def load(cls: Type[R], X_test: Any, y_test: Any, metrics: List[types.Metric], path: str = "report") -> R:
+        """Loads a report from an archive. See the save report class method.
+
+        The class method loads the stored models and results from the archive.
+        The testing data and metrics must be provided separately.
+
+        Args:
+            X_test (Any): Testing data. Should be identical to the original report.
+            y_test (Any): Ground truth prediction data. Should be identical to the original report.
+            metrics (List[types.Metric]): List of metrics. Should be identical to the original report.
+            path (str): Path to archived report. Don't append ".zip".
+                Defaults to 'report' since this is the default path for saving.
+
+        Returns:
+            R: Loaded report class
+        """
+        archive_type = "zip"
+        mode = "rb"
+        results_dir = "trackable_results"
+        models_sub_folder = "models"
+        zip_path = f"{path}.{archive_type}"
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            shutil.unpack_archive(zip_path, tmp_dir, format=archive_type)
+
+            models_dir = os.path.join(tmp_dir, models_sub_folder)
+            models = {}
+            results = []
+            for model_name in os.listdir(models_dir):
+                model_path = os.path.join(models_dir, model_name)
+                with open(model_path, mode) as f:
+                    loaded_model = pickle.load(f)
+                    models[model_name] = loaded_model
+
+                results_path = os.path.join(tmp_dir, results_dir)
+                with open(results_path, mode) as f:
+                    results = pickle.load(f)
+
+        report = cls(X_test, y_test, metrics)
+        report._models = models
+        report._results = results
+        return report
+
     def add_model(
         self,
         model: types.GenericModel,
@@ -115,3 +165,82 @@ def generate(self, highlight: Literal["max", "min", False] = "max") -> Union[Sty
         if highlight not in ("max", "min", False):
             raise TypeError("Highlight must be one of: max, min, None")
         return results
+
+    def get_model(self, name: str) -> types.GenericModel:
+        """Get a model from a report given its name.
+
+        Args:
+            name (str): Name of a model
+
+        Raises:
+            ModelDoesNotExistError: Raised if the given model name does not exist in the report
+
+        Returns:
+            types.GenericModel: A generic model
+        """
+        try:
+            return self._models[name]
+        except KeyError:
+            raise ModelDoesNotExistError(f"Model '{name}' does not exist.")
+
+    def remove_model(self, name: str) -> types.GenericModel:
+        """Removes a model given its name. The model is returned and then
+        removed from the report.
+
+        Args:
+            name (str): Name of a model
+
+        Raises:
+            ModelDoesNotExistError: Raised if the given model name does not exist in the report
+
+        Returns:
+            types.GenericModel: A generic model
+        """
+        try:
+            return self._models.pop(name)
+        except KeyError:
+            raise ModelDoesNotExistError(f"Model '{name}' does not exist.")
+
+    def save(self, path: str = "report") -> None:
+        """Saves report to an archive which can be loaded later.
+        Saved reports will not be overwritten and will raise an exception.
+
+        It is assumed that each model in the report can be pickled.
+
+        Note: X_test, y_test and metrics are NOT saved.
+        They should be saved separetely.
+
+        Args:
+            path (str, optional): Path to save archive.
+            Defaults to "report" in current working directory.
+            ".zip" is appended automatically.
+
+        Raises:
+            ArchiveAlreadyExistsError: Raised when a report archive with
+            the same name already exists.
+        """
+        archive = "zip"
+        mode = "wb"
+        results = "trackable_results"
+        models_sub_folder = "models"
+        zip_path = f"{path}.{archive}"
+
+        if os.path.isfile(zip_path):
+            raise ArchiveAlreadyExistsError(
+                f"{zip_path} already exists. Provide another path to avoid overwriting existing archives."
+            )
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            models_dir = os.path.join(tmp_dir, models_sub_folder)
+            os.mkdir(models_dir)
+
+            for name, model in self._models.items():
+                model_path = os.path.join(models_dir, name)
+                with open(model_path, mode) as f:
+                    pickle.dump(model, f)
+
+            results_path = os.path.join(tmp_dir, results)
+            with open(results_path, mode) as f:
+                pickle.dump(self._results, f)
+
+            shutil.make_archive(path, archive, tmp_dir)

trackable/types.py

@@ -11,4 +11,3 @@ class GenericModel(Protocol):
 
     def predict(self, data: Any) -> Any:
         """Generic predict method."""
-        ...