add type checking for dataframe setter

OpenCOMPES · Sep 22, 2023 · 7ece180 · 7ece180
1 parent 1bcde1a
commit 7ece180
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 0 deletions.
diff --git a/sed/core/processor.py b/sed/core/processor.py
@@ -173,6 +173,15 @@ def dataframe(self, dataframe: Union[pd.DataFrame, ddf.DataFrame]):
         Args:
             dataframe (Union[pd.DataFrame, ddf.DataFrame]): The dataframe object to set.
         """
+        if not isinstance(dataframe, (pd.DataFrame, ddf.DataFrame)) or not isinstance(
+            dataframe,
+            self._dataframe.__class__,
+        ):
+            raise ValueError(
+                "'dataframe' has to be a Pandas or Dask dataframe and has to be of the same kind "
+                "as the dataframe loaded into the SedProcessor!.\n"
+                f"Loaded type: {self._dataframe.__class__}, provided type: {dataframe}.",
+            )
         self._dataframe = dataframe
 
     @property

diff --git a/tests/test_processor.py b/tests/test_processor.py
@@ -155,6 +155,8 @@ def test_attributes_setters():
     processor.load(files=files, metadata={"test": {"key1": "value1"}})
     dataframe = processor.dataframe
     assert isinstance(dataframe, ddf.DataFrame)
+    with pytest.raises(ValueError):
+        processor.dataframe = dataframe["X"]
     dataframe["X"] = dataframe["Y"]
     processor.dataframe = dataframe
     np.testing.assert_allclose(