Allow dict as info field

pyprobe/cell.py

@@ -5,7 +5,7 @@
 import time
 import warnings
 import zipfile
-from typing import Callable, Dict, List, Optional
+from typing import Any, Callable, Dict, List, Optional
 
 import distinctipy
 import polars as pl
@@ -22,7 +22,7 @@
 class Cell(BaseModel):
     """A class for a cell in a battery experiment."""
 
-    info: Dict[str, Optional[str | int | float]]
+    info: Dict[str, Optional[str | int | float | Dict[Any, Any]]]
     """Dictionary containing information about the cell.
     The dictionary must contain a 'Name' field, other information may include
     channel number or other rig information.
@@ -32,8 +32,8 @@ class Cell(BaseModel):
 
     @field_validator("info")
     def check_and_set_name(
-        cls, info: Dict[str, Optional[str | int | float]]
-    ) -> Dict[str, Optional[str | int | float]]:
+        cls, info: Dict[str, Optional[str | int | float | Dict[Any, Any]]]
+    ) -> Dict[str, Optional[str | int | float | Dict[Any, Any]]]:
         """Validate the `info` field.
 
         Checks that a `Name` field is present in the `info` dictionary, if not it is
@@ -240,7 +240,7 @@ def _write_parquet(
 
     @staticmethod
     def _get_filename(
-        info: Dict[str, Optional[str | int | float]],
+        info: Dict[str, Optional[str | int | float | Dict[Any, Any]]],
         filename_function: Callable[[str], str],
         filename_inputs: List[str],
     ) -> str:

pyprobe/filters.py

@@ -4,10 +4,9 @@
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast
 
 import polars as pl
-from pydantic import Field
 
 from pyprobe import utils
-from pyprobe.rawdata import RawData, default_column_definitions
+from pyprobe.rawdata import RawData
 
 if TYPE_CHECKING:
     from pyprobe.typing import (  # , FilterToStepType
@@ -277,14 +276,14 @@ class Procedure(RawData):
     """A class for a procedure in a battery experiment."""
 
     readme_dict: Dict[str, Dict[str, List[str | int | Tuple[int, int, int]]]]
+    """A dictionary representing the data contained in the README yaml file."""
 
-    base_dataframe: pl.LazyFrame | pl.DataFrame
-    info: Dict[str, Optional[str | int | float]]
-    column_definitions: Dict[str, str] = Field(
-        default_factory=lambda: default_column_definitions.copy()
-    )
-    step_descriptions: Dict[str, List[Optional[str | int]]] = {}
     cycle_info: List[Tuple[int, int, int]] = []
+    """A list of tuples representing the cycle information from the README yaml file.
+
+    The tuple format is
+    :code:`(start step (inclusive), end step (inclusive), cycle count)`.
+    """
 
     def model_post_init(self, __context: Any) -> None:
         """Create a procedure class."""
@@ -449,12 +448,12 @@ def load_external_file(self, filepath: str) -> pl.LazyFrame:
 class Experiment(RawData):
     """A class for an experiment in a battery experimental procedure."""
 
-    base_dataframe: pl.LazyFrame | pl.DataFrame
-    info: Dict[str, Optional[str | int | float]]
-    column_definitions: Dict[str, str] = Field(
-        default_factory=lambda: default_column_definitions.copy()
-    )
-    cycle_info: List[Tuple[int, int, int]]
+    cycle_info: List[Tuple[int, int, int]] = []
+    """A list of tuples representing the cycle information from the README yaml file.
+
+    The tuple format is
+    :code:`(start step (inclusive), end step (inclusive), cycle count)`.
+    """
 
     def model_post_init(self, __context: Any) -> None:
         """Create an experiment class."""
@@ -483,12 +482,12 @@ def model_post_init(self, __context: Any) -> None:
 class Cycle(RawData):
     """A class for a cycle in a battery experimental procedure."""
 
-    base_dataframe: pl.LazyFrame | pl.DataFrame
-    info: Dict[str, Optional[str | int | float]]
-    column_definitions: Dict[str, str] = Field(
-        default_factory=lambda: default_column_definitions.copy()
-    )
-    cycle_info: List[Tuple[int, int, int]]
+    cycle_info: List[Tuple[int, int, int]] = []
+    """A list of tuples representing the cycle information from the README yaml file.
+
+    The tuple format is
+    :code:`(start step (inclusive), end step (inclusive), cycle count)`.
+    """
 
     def model_post_init(self, __context: Any) -> None:
         """Create a cycle class."""
@@ -516,12 +515,6 @@ def model_post_init(self, __context: Any) -> None:
 class Step(RawData):
     """A class for a step in a battery experimental procedure."""
 
-    base_dataframe: pl.LazyFrame | pl.DataFrame
-    info: Dict[str, Optional[str | int | float]]
-    column_definitions: Dict[str, str] = Field(
-        default_factory=lambda: default_column_definitions.copy()
-    )
-
     def model_post_init(self, __context: Any) -> None:
         """Create a step class."""
         self.zero_column(

pyprobe/rawdata.py

@@ -50,12 +50,15 @@ class RawData(Result):
     This defines the PyProBE format.
     """
 
-    base_dataframe: pl.LazyFrame | pl.DataFrame
-    info: Dict[str, Optional[str | int | float]]
     column_definitions: Dict[str, str] = Field(
         default_factory=lambda: default_column_definitions.copy()
     )
     step_descriptions: Dict[str, List[Optional[str | int]]] = {}
+    """A dictionary containing the fields 'Step' and 'Description'.
+
+    - 'Step' is a list of step numbers.
+    - 'Description' is a list of corresponding descriptions in PyBaMM Experiment format.
+    """
 
     @field_validator("base_dataframe")
     @classmethod

pyprobe/result.py

@@ -35,7 +35,7 @@ class Config:
 
     base_dataframe: Union[pl.LazyFrame, pl.DataFrame]
     """The data as a polars DataFrame or LazyFrame."""
-    info: Dict[str, Optional[str | int | float]]
+    info: Dict[str, Optional[str | int | float | Dict[Any, Any]]]
     """Dictionary containing information about the cell."""
     column_definitions: Dict[str, str] = Field(default_factory=dict)
     """A dictionary containing the definitions of the columns in the data."""
@@ -364,7 +364,7 @@ def build(
                 | Dict[str, NDArray[np.float64] | List[float]]
             ]
         ],
-        info: Dict[str, Optional[str | int | float]],
+        info: Dict[str, Optional[str | int | float | Dict[Any, Any]]],
     ) -> "Result":
         """Build a Result object from a list of dataframes.