improve BLData type hints

bletl/types.py

@@ -32,7 +32,73 @@
     """Additions from pipetting."""
 
 
-class BLData(dict):
+class FilterTimeSeries:
+    """Generalizable data type for calibrated timeseries."""
+
+    @property
+    def wells(self) -> typing.Tuple[str, ...]:
+        """Well IDs that were measured."""
+        return tuple(self.time.columns)
+
+    def __init__(self, time_df: pandas.DataFrame, value_df: pandas.DataFrame):
+        self.time = time_df
+        self.value = value_df
+
+    def get_timeseries(
+        self, well: str, *, last_cycle: Optional[int] = None
+    ) -> Tuple[numpy.ndarray, numpy.ndarray]:
+        """Retrieves (time, value) for a specific well.
+
+        Parameters
+        ----------
+        well : str
+            Well id to retrieve.
+        last_cycle : int, optional
+            Cycle number of the last cycle to be included (defaults to all cycles).
+
+        Returns
+        -------
+        x : numpy.ndarray
+            Timepoints of measurements.
+        y : numpy.ndarray
+            Measured values.
+        """
+        if last_cycle is not None and last_cycle <= 0:
+            raise ValueError(f"last_cycle must be > 0")
+        x = numpy.array(self.time[well])[:last_cycle]
+        y = numpy.array(self.value[well])[:last_cycle]
+        return x, y
+
+    def get_unified_dataframe(self, well: Optional[str] = None) -> pandas.DataFrame:
+        """Retrieves a DataFrame with unified time on index.
+
+        Parameters
+        ----------
+        well : str, optional
+            Well id from which time is taken.
+            If `None`, the first well is used.
+
+        Returns
+        -------
+        unified_df : pandas.DataFrame
+            Dataframe with unified time on index.
+        """
+        if not well is None:
+            if not well in self.time.columns:
+                raise KeyError("Could not find well id")
+            time = self.time.loc[:, well]
+        else:
+            time = self.time.iloc[:, 0]
+
+        new_index = pandas.Index(time, name="time in h")
+        unified_df = self.value.set_index(new_index)
+        return unified_df
+
+    def __repr__(self):
+        return f"FilterTimeSeries({len(self.time)} cycles, {len(self.time.columns)} wells)"
+
+
+class BLData(Dict[str, FilterTimeSeries]):
     """Standardized data type for BioLector data."""
 
     def __init__(
@@ -223,72 +289,6 @@
         )
 
 
-class FilterTimeSeries:
-    """Generalizable data type for calibrated timeseries."""
-
-    @property
-    def wells(self) -> typing.Tuple[str, ...]:
-        """Well IDs that were measured."""
-        return tuple(self.time.columns)
-
-    def __init__(self, time_df: pandas.DataFrame, value_df: pandas.DataFrame):
-        self.time = time_df
-        self.value = value_df
-
-    def get_timeseries(
-        self, well: str, *, last_cycle: Optional[int] = None
-    ) -> Tuple[numpy.ndarray, numpy.ndarray]:
-        """Retrieves (time, value) for a specific well.
-
-        Parameters
-        ----------
-        well : str
-            Well id to retrieve.
-        last_cycle : int, optional
-            Cycle number of the last cycle to be included (defaults to all cycles).
-
-        Returns
-        -------
-        x : numpy.ndarray
-            Timepoints of measurements.
-        y : numpy.ndarray
-            Measured values.
-        """
-        if last_cycle is not None and last_cycle <= 0:
-            raise ValueError(f"last_cycle must be > 0")
-        x = numpy.array(self.time[well])[:last_cycle]
-        y = numpy.array(self.value[well])[:last_cycle]
-        return x, y
-
-    def get_unified_dataframe(self, well: Optional[str] = None) -> pandas.DataFrame:
-        """Retrieves a DataFrame with unified time on index.
-
-        Parameters
-        ----------
-        well : str, optional
-            Well id from which time is taken.
-            If `None`, the first well is used.
-
-        Returns
-        -------
-        unified_df : pandas.DataFrame
-            Dataframe with unified time on index.
-        """
-        if not well is None:
-            if not well in self.time.columns:
-                raise KeyError("Could not find well id")
-            time = self.time.loc[:, well]
-        else:
-            time = self.time.iloc[:, 0]
-
-        new_index = pandas.Index(time, name="time in h")
-        unified_df = self.value.set_index(new_index)
-        return unified_df
-
-    def __repr__(self):
-        return f"FilterTimeSeries({len(self.time)} cycles, {len(self.time.columns)} wells)"
-
-
 class BLDParser:
     """Abstract type for parsers that read BioLector CSV files."""
 

tests/test_core.py

@@ -319,6 +319,25 @@ def test_NoMeasurements_Warning(self):
         with pytest.warns(NoMeasurementData):
             bletl.parse(file_with_no_measurements)
 
+    def test_get_unified_dataframe_well_selection(self):
+        # Create test data
+        time_df = pandas.DataFrame({"A01": [0.0, 1.0, 2.0], "A02": [0.0, 1.0, 2.0]})
+        value_df = pandas.DataFrame({"A01": [1.0, 2.0, 3.0], "A02": [1.5, 2.5, 3.5]})
+
+        fts = bletl.FilterTimeSeries(time_df, value_df)
+
+        # Test with valid well ID - should return DataFrame with correct time values
+        df = fts.get_unified_dataframe(well="A01")
+        numpy.testing.assert_array_equal(df.index.values, [0.0, 1.0, 2.0])
+
+        # Test with non-existent well ID - should raise KeyError
+        with pytest.raises(KeyError, match="Could not find well id"):
+            fts.get_unified_dataframe(well="X99")
+
+        # Test with None
+        df_default = fts.get_unified_dataframe(well=None)
+        numpy.testing.assert_array_equal(df_default.index.values, [0.0, 1.0, 2.0])
+
 
 class TestBL1Calibration:
     def test_calibration_data_type(self):