Replace filling missing columns with ignoring file and providing a warning

Author: Tom Holland

pyprobe/cyclers/biologic.py

@@ -3,6 +3,7 @@
 
 import glob
 import re
+import warnings
 from datetime import datetime
 from typing import List
 
@@ -72,7 +73,6 @@ def read_file(filepath: str) -> pl.DataFrame:
         dataframe = dataframe.with_columns(
             (pl.col("time/s") * 1000000 + start).cast(pl.Datetime).alias("Date")
         )
-
         return dataframe
 
     @classmethod
@@ -110,32 +110,24 @@ def raw_dataframe(self) -> pl.DataFrame:
         files = glob.glob(self.input_data_path)
         files = self.sort_files(files)
         dataframes = [self.read_file(file) for file in files]
-        all_columns = [col for df in dataframes for col in df.columns]
+        all_columns = set([col for df in dataframes for col in df.columns])
+        indices_to_remove = []
         for i in range(len(dataframes)):
-            dataframes[i] = self.fill_missing_columns(dataframes[i], all_columns)
+            if len(dataframes[i].columns) < len(all_columns):
+                indices_to_remove.append(i)
+                warnings.warn(
+                    f"File {files[i]} has missing columns, it has not been read."
+                )
+                continue
             if i > 0:
                 dataframes[i] = dataframes[i].with_columns(
                     pl.col("Ns") + dataframes[i - 1]["Ns"].max() + 1
                 )
+        dataframes = [
+            df for i, df in enumerate(dataframes) if i not in indices_to_remove
+        ]
         return pl.concat(dataframes, how="vertical")
 
-    def fill_missing_columns(
-        self, dataframe: pl.DataFrame, column_list: List[str]
-    ) -> pl.DataFrame:
-        """Fill missing columns in a DataFrame with zeros.
-
-        Args:
-            dataframe: The DataFrame to fill.
-            column_list: The list of required columns.
-
-        Returns:
-            pl.DataFrame: The DataFrame with filled columns.
-        """
-        missing_columns = set(column_list) - set(dataframe.columns)
-        for col in missing_columns:
-            dataframe = dataframe.with_columns(pl.zeros(dataframe.height).alias(col))
-        return dataframe
-
     @property
     def processed_dataframe(self) -> pl.DataFrame:
         """Process a DataFrame from battery cycler data.

tests/cyclers/test_biologic.py

@@ -28,40 +28,6 @@ def test_read_file(biologic_cycler):
         "Q charge/mA.h",
         "Date",
     ]
-    assert 0 == 1
-
-
-def test_fill_missing_columns(biologic_cycler):
-    """Test the fill_missing_columns method."""
-    dataframe = pl.DataFrame(
-        {
-            "time/s": [0.0, 1.0, 2.0, 3.0],
-            "Ns": [0, 1, 2, 3],
-            "Ecell/V": [4, 5, 6, 7],
-        }
-    )
-    all_columns = [
-        "time/s",
-        "Ns",
-        "Ecell/V",
-        "I/mA",
-        "Q charge/mA.h",
-        "Q discharge/mA.h",
-    ]
-    filled_dataframe = biologic_cycler.fill_missing_columns(dataframe, all_columns)
-    expected_dataframe = pl.DataFrame(
-        {
-            "time/s": [0.0, 1.0, 2.0, 3.0],
-            "Ns": [0, 1, 2, 3],
-            "Ecell/V": [4, 5, 6, 7],
-            "I/mA": [0.0, 0.0, 0.0, 0.0],
-            "Q charge/mA.h": [0.0, 0.0, 0.0, 0.0],
-            "Q discharge/mA.h": [0.0, 0.0, 0.0, 0.0],
-        }
-    )
-    pl_testing.assert_frame_equal(
-        filled_dataframe, expected_dataframe, check_column_order=False
-    )
 
 
 def test_sort_files(biologic_cycler):