ionworks-api 0.1.0__py3-none-any.whl → 0.1.3__py3-none-any.whl

ionworks/__init__.py

@@ -1,3 +1,28 @@
-from .client import Ionworks, IonworksError
+"""
+Ionworks API Client.
 
-__all__ = ["Ionworks", "IonworksError"]
+A Python client for interacting with the Ionworks platform for battery cell
+testing, simulation, and modeling.
+
+Example:
+-------
+>>> from ionworks import Ionworks
+>>> client = Ionworks()
+>>> specs = client.cell_spec.list()
+"""
+
+from .client import Ionworks
+from .errors import IonworksError
+from .validators import (
+    MeasurementValidationError,
+    get_dataframe_backend,
+    set_dataframe_backend,
+)
+
+__all__ = [
+    "Ionworks",
+    "IonworksError",
+    "MeasurementValidationError",
+    "get_dataframe_backend",
+    "set_dataframe_backend",
+]

ionworks/cell_instance.py

@@ -1,3 +1,11 @@
+"""
+Cell instance client for managing individual cell records.
+
+This module provides the :class:`CellInstanceClient` for creating, reading,
+updating, and deleting cell instances, which represent individual physical
+battery cells linked to a cell specification.
+"""
+
 from typing import Any
 
 from pydantic import ValidationError
@@ -12,13 +20,20 @@ from .models import (
 
 
 class CellInstanceClient:
-    def __init__(self, client: Any):
+    """Client for managing cell instances."""
+
+    def __init__(self, client: Any) -> None:
+        """Initialize the CellInstanceClient.
+
+        Parameters
+        ----------
+        client : Any
+            The HTTP client instance used for API calls.
+        """
         self.client = client
 
     def get(self, cell_instance_id: str) -> CellInstance:
-        """
-        Get a specific cell instance by ID.
-        """
+        """Get a specific cell instance by ID."""
         endpoint = f"/cell_instances/{cell_instance_id}"
         try:
             response_data = self.client.get(endpoint)
@@ -30,9 +45,7 @@ class CellInstanceClient:
             raise RuntimeError(f"API call to {endpoint} failed: {e}") from e
 
     def get_by_slug(self, cell_instance_slug: str) -> CellInstance:
-        """
-        Get a specific cell instance by slug.
-        """
+        """Get a specific cell instance by slug."""
         endpoint = f"/cell_instances/slug/{cell_instance_slug}"
         try:
             response_data = self.client.get(endpoint)
@@ -43,9 +56,7 @@ class CellInstanceClient:
             raise RuntimeError(f"API call to {endpoint} failed: {e}") from e
 
     def list(self, cell_spec_id: str) -> list[CellInstance]:
-        """
-        List all cell instances for a cell specification by specification ID.
-        """
+        """List all cell instances for a cell specification by specification ID."""
         endpoint = f"/cell_specifications/{cell_spec_id}/cell_instances"
         try:
             response_data = self.client.get(endpoint)
@@ -74,14 +85,13 @@ class CellInstanceClient:
     #     return [CellMeasurementDetailBase(**item) for item in response_data]
 
     def update(self, cell_instance_id: str, data: dict[str, Any]) -> CellInstance:
-        """
-        Update an existing cell instance.
+        """Update an existing cell instance.
 
         Parameters
         ----------
         cell_instance_id : str
             The ID of the cell instance to update.
-        data : dict
+        data : dict[str, Any]
             Dictionary containing the fields to update.
 
         Returns
@@ -94,30 +104,36 @@ class CellInstanceClient:
         return CellInstance(**response_data)
 
     def delete(self, cell_instance_id: str) -> None:
-        """
-        Delete a cell instance by ID.
-        """
+        """Delete a cell instance by ID."""
         endpoint = f"/cell_instances/{cell_instance_id}"
         self.client.delete(endpoint)
 
     # Renamed and moved method from Data class
     def detail(self, cell_instance_id: str) -> CellInstanceDetail:
-        """
-        Loads the full details for a cell instance from a single API endpoint,
-        including its spec, instance data, steps DataFrame, and time series DataFrame.
-        Uses the /detail endpoint.
+        """Load the full details for a cell instance.
 
-        Args:
-            organization: The organization slug.
-            cell_instance_id: The cell instance ID.
+        Loads the full details for a cell instance from a single API endpoint,
+        including its spec, instance data, steps DataFrame, and time series
+        DataFrame. Uses the /detail endpoint.
 
-        Returns:
-            A FullCellInstance object containing validated data.
+        Parameters
+        ----------
+        cell_instance_id : str
+            The cell instance ID.
 
-        Raises:
-            RuntimeError: If the API call fails.
-            ValueError: If the API response format is unexpected or parsing fails.
-            ValidationError: If the response data doesn't match expected models.
+        Returns
+        -------
+        CellInstanceDetail
+            Object containing validated cell instance data.
+
+        Raises
+        ------
+        RuntimeError
+            If the API call fails.
+        ValueError
+            If the API response format is unexpected or parsing fails.
+        ValidationError
+            If the response data doesn't match expected models.
         """
         endpoint = f"/cell_instances/{cell_instance_id}/detail"
         try:
@@ -127,9 +143,9 @@ class CellInstanceClient:
         except ValidationError as e:
             # Pydantic validation error during parsing
             raise ValueError(f"Response validation failed for {endpoint}: {e}") from e
-        except ValueError as e:
+        except ValueError:
             # Re-raise our explicit format validation errors
-            raise e
+            raise
         except Exception as e:
             # Catch client errors or any other unexpected issues
             raise RuntimeError(
@@ -137,19 +153,13 @@ class CellInstanceClient:
             ) from e
 
     def create(self, cell_spec_id: str, data: dict[str, Any]) -> CellInstance:
-        """
-        Create a new cell instance under a cell specification by
-        specification ID.
-        """
+        """Create a new cell instance under a cell specification by specification ID."""
         endpoint = f"/cell_specifications/{cell_spec_id}/cell_instances"
         response_data = self.client.post(endpoint, data)
         return CellInstance(**response_data)
 
     def create_or_get(self, cell_spec_id: str, data: dict[str, Any]) -> CellInstance:
-        """
-        Create a new cell instance if it doesn't exist, otherwise get the
-        existing one.
-        """
+        """Create a new cell instance if it doesn't exist, otherwise get it."""
         try:
             return self.create(cell_spec_id, data)
         except IonworksError as e:
@@ -157,11 +167,11 @@ class CellInstanceClient:
                 existing_id = e.data.get("existing_cell_instance_id")
                 if existing_id:
                     return self.get(existing_id)
-            raise e
+            raise
 
     def list_with_measurements(self, cell_spec_id: str) -> CellSpecificationInstances:
-        """
-        List all instances and measurements for a cell specification.
+        """List all instances and measurements for a cell specification.
+
         Returns normalized data: separate lists for specification, instances, and
         measurements.
         """
@@ -179,8 +189,8 @@ class CellInstanceClient:
     def list_with_measurements_and_steps(
         self, cell_spec_id: str
     ) -> CellSpecificationInstancesWithSteps:
-        """
-        List all instances, measurements, and steps for a cell specification.
+        """List all instances, measurements, and steps for a cell specification.
+
         Returns normalized data: separate lists for specification, instances,
         measurements, and steps.
         """
@@ -195,7 +205,3 @@ class CellInstanceClient:
             raise ValueError(f"Response validation failed for {endpoint}: {e}") from e
         except Exception as e:
             raise RuntimeError(f"API call to {endpoint} failed: {e}") from e
-
-
-class CellMeasurementClient:
-    pass

ionworks/cell_measurement.py

@@ -1,11 +1,18 @@
+"""Cell measurement client for managing time series test data.
+
+This module provides the :class:`CellMeasurementClient` for uploading,
+retrieving, and managing measurement data from battery cell testing. It
+supports efficient upload of large datasets using signed URLs and parquet
+format.
+"""
+
 from __future__ import annotations
 
 import io
 from typing import Any
 
 import pandas as pd
-import pyarrow as pa
-import pyarrow.parquet as pq
+import polars as pl
 from pydantic import ValidationError
 import requests
 
@@ -18,46 +25,73 @@ from .models import (
     ConfirmUploadResponse,
     InitiateUploadResponse,
 )
-from .validators import df_to_dict_validator
+from .validators import DataFrame, df_to_dict_validator, validate_measurement_data
+
+
+def _to_polars(df: DataFrame | dict) -> pl.DataFrame:
+    """Convert pandas DataFrame or dict to polars DataFrame.
+
+    Parameters
+    ----------
+    df : DataFrame | dict
+        Input data as pandas DataFrame, polars DataFrame, or dict.
+
+    Returns
+    -------
+    pl.DataFrame
+        Polars DataFrame.
+    """
+    if isinstance(df, pl.DataFrame):
+        return df
+    if isinstance(df, pd.DataFrame):
+        return pl.from_pandas(df)
+    if isinstance(df, dict):
+        return pl.DataFrame(df)
+    raise TypeError(f"Expected DataFrame or dict, got {type(df).__name__}")
 
 
 class CellMeasurementClient:
-    def __init__(self, client: Any):
+    """Client for managing cell measurement data."""
+
+    #: Default timeout for signed URL uploads as (connect, read) in seconds.
+    UPLOAD_TIMEOUT: tuple[float, float] = (10, 300)
+
+    def __init__(self, client: Any) -> None:
+        """Initialize the CellMeasurementClient.
+
+        Parameters
+        ----------
+        client : Any
+            The HTTP client instance for making API calls.
+        """
         self.client = client
 
     def list(self, cell_instance_id: str) -> list[CellMeasurement]:
-        """
-        List all cell measurements for a cell instance by instance ID.
-        """
+        """List all cell measurements for a cell instance by instance ID."""
         endpoint = f"/cell_instances/{cell_instance_id}/cell_measurements"
         response_data = self.client.get(endpoint)
         return [CellMeasurement(**item) for item in response_data]
 
     def get(self, measurement_id: str) -> CellMeasurement:
-        """
-        Get a specific cell measurement by its ID only.
-        """
+        """Get a specific cell measurement by its ID only."""
         endpoint = f"/cell_measurements/{measurement_id}"
         response_data = self.client.get(endpoint)
         return CellMeasurement(**response_data)
 
     def detail(self, measurement_id: str) -> CellMeasurementDetail:
-        """
-        Get the full details for a cell measurement by measurement ID only.
-        """
+        """Get the full details for a cell measurement by measurement ID only."""
         endpoint = f"/cell_measurements/{measurement_id}/detail"
         response_data = self.client.get(endpoint)
         return CellMeasurementDetail(**response_data)
 
     def update(self, measurement_id: str, data: dict[str, Any]) -> CellMeasurement:
-        """
-        Update an existing cell measurement.
+        """Update an existing cell measurement.
 
         Parameters
         ----------
         measurement_id : str
             The ID of the cell measurement to update.
-        data : dict
+        data : dict[str, Any]
             Dictionary containing the fields to update.
 
         Returns
@@ -70,17 +104,15 @@ class CellMeasurementClient:
         return CellMeasurement(**response_data)
 
     def delete(self, measurement_id: str) -> None:
-        """
-        Delete a cell measurement by measurement ID only.
-        """
+        """Delete a cell measurement by measurement ID only."""
         endpoint = f"/cell_measurements/{measurement_id}"
         self.client.delete(endpoint)
 
     def list_with_steps(
         self, cell_instance_id: str
     ) -> CellInstanceMeasurementsWithSteps:
-        """
-        List all measurements and steps for a specific cell instance.
+        """List all measurements and steps for a specific cell instance.
+
         Returns normalized data: separate lists for instance, measurements, and
         steps.
         """
@@ -96,8 +128,7 @@ class CellMeasurementClient:
     def get_by_slugs(
         self, cell_instance_slug: str, measurement_slug: str
     ) -> CellMeasurement:
-        """
-        Get a specific cell measurement by instance slug and measurement slug.
+        """Get a specific cell measurement by instance slug and measurement slug.
 
         Returns metadata only (no steps or time series).
         """
@@ -108,23 +139,21 @@ class CellMeasurementClient:
         response_data = self.client.get(endpoint)
         return CellMeasurement(**response_data)
 
-    def _dataframe_to_parquet(self, df: pd.DataFrame) -> bytes:
-        """
-        Convert a pandas DataFrame to parquet bytes.
+    def _dataframe_to_parquet(self, df: pl.DataFrame) -> bytes:
+        """Convert a polars DataFrame to parquet bytes.
 
         Parameters
         ----------
-        df : pd.DataFrame
+        df : pl.DataFrame
             The DataFrame to convert.
 
         Returns
         -------
         bytes
-            Parquet file content (uses built-in zstd compression).
+            Parquet file content (uses zstd compression).
         """
-        table = pa.Table.from_pandas(df)
         buffer = io.BytesIO()
-        pq.write_table(table, buffer, compression="zstd")
+        df.write_parquet(buffer, compression="zstd")
         return buffer.getvalue()
 
     def _initiate_upload(
@@ -133,8 +162,7 @@ class CellMeasurementClient:
         name: str,
         notes: str | None = None,
     ) -> InitiateUploadResponse:
-        """
-        Initiate a signed URL upload for a new measurement.
+        """Initiate a signed URL upload for a new measurement.
 
         Parameters
         ----------
@@ -142,7 +170,7 @@ class CellMeasurementClient:
             The ID of the cell instance.
         name : str
             Name for the new measurement.
-        notes : str, optional
+        notes : str | None
             Optional notes for the measurement.
 
         Returns
@@ -162,8 +190,7 @@ class CellMeasurementClient:
         return InitiateUploadResponse(**response_data)
 
     def _upload_to_signed_url(self, signed_url: str, parquet_bytes: bytes) -> None:
-        """
-        Upload parquet bytes to a signed URL.
+        """Upload parquet bytes to a signed URL.
 
         Parameters
         ----------
@@ -185,6 +212,7 @@ class CellMeasurementClient:
                 url,
                 data=parquet_bytes,
                 headers={"Content-Type": "application/octet-stream"},
+                timeout=self.UPLOAD_TIMEOUT,
             )
             response.raise_for_status()
         except requests.exceptions.RequestException as e:
@@ -197,8 +225,7 @@ class CellMeasurementClient:
         measurement_data: dict[str, Any],
         steps: dict[str, list[Any]] | None = None,
     ) -> ConfirmUploadResponse:
-        """
-        Confirm a signed URL upload after successful file upload.
+        """Confirm a signed URL upload after successful file upload.
 
         This creates the measurement record in the database. The measurement
         data must be passed again (same as initiate) to ensure no orphaned
@@ -210,10 +237,12 @@ class CellMeasurementClient:
             The pre-generated ID for the measurement (from initiate).
         cell_instance_id : str
             The ID of the parent cell instance.
-        measurement_data : dict
-            Measurement metadata (name, notes, etc.) - same as passed to initiate.
-        steps : dict, optional
-            Pre-calculated steps data. If not provided, will be auto-calculated.
+        measurement_data : dict[str, Any]
+            Measurement metadata (name, notes, etc.) — same as passed to
+            initiate.
+        steps : dict[str, list[Any]] | None
+            Pre-calculated steps data. If not provided, will be
+            auto-calculated.
 
         Returns
         -------
@@ -237,9 +266,7 @@ class CellMeasurementClient:
         cell_instance_id: str,
         measurement_detail: dict[str, Any],
     ) -> CellMeasurementBundleResponse:
-        """
-        Create a new cell measurement with steps and time series data
-        for a cell instance.
+        """Create a new cell measurement with steps and time series data.
 
         Uses signed URL upload for better performance with large datasets.
         Data is uploaded directly to storage as parquet, bypassing backend
@@ -250,9 +277,10 @@ class CellMeasurementClient:
         ----------
         cell_instance_id : str
             The ID of the cell instance to create the measurement for.
-        measurement_detail : dict
+        measurement_detail : dict[str, Any]
             Dictionary containing 'measurement', 'steps', and 'time_series'.
-            - measurement: dict with 'name' (required) and 'notes' (optional)
+
+            - measurement: dict with 'name' (required) and 'notes'
             - time_series: pandas DataFrame or dict with time series data
             - steps: optional dict with pre-calculated steps data
 
@@ -261,24 +289,35 @@ class CellMeasurementClient:
         CellMeasurementBundleResponse
             Response containing the created measurement, steps count, and
             file path.
+
+        Raises
+        ------
+        MeasurementValidationError
+            If data validation fails. Validation checks include:
+            - Positive current should correspond to discharge
+            - Cumulative values should reset at each step
         """
         measurement_info = measurement_detail["measurement"]
         name = measurement_info["name"]
         notes = measurement_info.get("notes")
 
-        # Step 1: Initiate upload (validates metadata, returns signed URL, no DB record)
+        # Step 1: Convert time_series to polars DataFrame
+        # Accepts pandas DataFrame, polars DataFrame, or dict
+        time_series = _to_polars(measurement_detail["time_series"])
+
+        # Step 2: Validate the data before any upload
+        validate_measurement_data(time_series)
+
+        # Step 3: Initiate upload (validates metadata, returns signed URL, no DB record)
         initiate_result = self._initiate_upload(cell_instance_id, name, notes)
 
-        # Step 2: Convert time_series to DataFrame if needed, then to parquet
-        time_series = measurement_detail["time_series"]
-        if isinstance(time_series, dict):
-            time_series = pd.DataFrame(time_series)
+        # Step 4: Convert to parquet
         parquet_bytes = self._dataframe_to_parquet(time_series)
 
-        # Step 3: Upload to signed URL
+        # Step 5: Upload to signed URL
         self._upload_to_signed_url(initiate_result.signed_url, parquet_bytes)
 
-        # Step 4: Confirm upload (creates the measurement record)
+        # Step 6: Confirm upload (creates the measurement record)
         steps = measurement_detail.get("steps")
         if steps is not None:
             # Convert DataFrame to dict if needed
@@ -301,15 +340,13 @@ class CellMeasurementClient:
         cell_instance_id: str,
         measurement_detail: dict[str, Any],
     ) -> CellMeasurementBundleResponse | CellMeasurement:
-        """
-        Create a new cell measurement if it doesn't exist, otherwise get the
-        existing one.
+        """Create a new cell measurement or get existing one.
 
         Parameters
         ----------
         cell_instance_id : str
             The ID of the cell instance to create the measurement for.
-        measurement_detail : dict
+        measurement_detail : dict[str, Any]
             Dictionary containing 'measurement', 'steps', and 'time_series'.
 
         Returns
@@ -335,5 +372,5 @@ class CellMeasurementClient:
                 raise ValueError(
                     f"Measurement '{measurement_name}' reported as duplicate "
                     "but could not be found"
-                )
+                ) from e
             raise