fix(datasets): improve save_dataframe function signature and implemen…

…tation

src/hyfi/utils/datasets/save.py

@@ -23,6 +23,31 @@ class DSSave:
     def __init__(self):
         pass
 
+    @staticmethod
+    def save_dataframe(
+        data: pd.DataFrame,
+        data_file: str,
+        data_dir: Optional[str] = None,
+        columns: Optional[Sequence[str]] = None,
+        index: bool = False,
+        filetype: Optional[str] = "parquet",
+        suffix: Optional[str] = None,
+        verbose: bool = False,
+        **kwargs,
+    ):
+        """Save data to a file"""
+        DSSave.save_dataframes(
+            data,
+            data_file,
+            data_dir=data_dir,
+            columns=columns,
+            index=index,
+            filetype=filetype,
+            suffix=suffix,
+            verbose=verbose,
+            **kwargs,
+        )
+
     @staticmethod
     def save_dataframes(
         data: Union[pd.DataFrame, dict],

tests/hyfi/utils/test_datasets_load.py

@@ -1,5 +1,6 @@
 import unittest
 from hyfi.utils.datasets.load import DSLoad
+from hyfi.utils.datasets.save import DSSave
 import pandas as pd
 
 
@@ -29,6 +30,17 @@ def test_load_dataframe_with_columns(self):
         self.assertIsInstance(dataframe, pd.DataFrame)
         self.assertListEqual(list(dataframe.columns), columns)
 
+    def test_save_dataframe(self):
+        # Test saving a dataframe
+        data = pd.DataFrame({"column1": [1, 2, 3], "column2": [4, 5, 6]})
+        data_file = "saved-data.csv"
+        data_dir = "tests/data"
+        DSSave.save_dataframe(data, data_file, data_dir=data_dir)
+        dataframe = DSLoad.load_dataframe(data_file, data_dir=data_dir)
+        self.assertIsNotNone(dataframe)
+        self.assertIsInstance(dataframe, pd.DataFrame)
+        self.assertListEqual(list(dataframe.columns), list(data.columns))
+
 
 if __name__ == "__main__":
     unittest.main()