test coverage for no dataset saved

cleanlab · Jan 12, 2025 · d611219 · d611219
1 parent 5c11295
commit d611219
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 3 deletions.
diff --git a/cleanlab/datalab/datalab.py b/cleanlab/datalab/datalab.py
@@ -592,7 +592,7 @@ def list_default_issue_types(self) -> List[str]:
             default_issue_types.extend(DEFAULT_CLEANVISION_ISSUES.keys())
         return default_issue_types
 
-    def save(self, path: str, force: bool = False) -> None:
+    def save(self, path: str, force: bool = False, save_dataset: bool = True) -> None:
         """Saves this Datalab object to file (all files are in folder at `path/`).
         We do not guarantee saved Datalab can be loaded from future versions of cleanlab.
 
@@ -608,7 +608,7 @@ def save(self, path: str, force: bool = False) -> None:
         ----
         You have to save the Dataset yourself separately if you want it saved to file.
         """
-        _Serializer.serialize(path=path, datalab=self, force=force)
+        _Serializer.serialize(path=path, datalab=self, force=force, save_dataset=save_dataset)
         save_message = f"Saved Datalab to folder: {path}"
         print(save_message)
 

diff --git a/cleanlab/datalab/internal/serialize.py b/cleanlab/datalab/internal/serialize.py
@@ -76,7 +76,7 @@ def _validate_version(datalab: Datalab) -> None:
             )
 
     @classmethod
-    def serialize(cls, path: str, datalab: Datalab, force: bool, save_dataset: bool = True) -> None:
+    def serialize(cls, path: str, datalab: Datalab, force: bool, save_dataset: bool) -> None:
         """Serializes the datalab object to disk.
 
         Parameters

diff --git a/tests/datalab/datalab/test_datalab.py b/tests/datalab/datalab/test_datalab.py
@@ -451,6 +451,14 @@ def test_save(self, lab, tmp_path, mock_issues, mock_issue_summary, monkeypatch)
     def test_load(self, lab, tmp_path, dataset, mock_issues, mock_issue_summary, monkeypatch):
         """Test that the save and load methods work."""
 
+        # Load without saved datalab raises warning
+        lab.save(tmp_path, force=True, save_dataset=False)
+        with pytest.warns(
+            UserWarning,
+            match="The Datalab object was loaded without a dataset and will have limited functionality.",
+        ):
+            loaded_lab = Datalab.load(tmp_path)
+
         monkeypatch.setattr(lab, "issues", mock_issues)
         monkeypatch.setattr(lab, "issue_summary", mock_issue_summary)