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

ionworks/cell_specification.py

@@ -1,3 +1,10 @@
+"""Cell specification client for managing cell type definitions.
+
+This module provides the :class:`CellSpecificationClient` for creating,
+reading, updating, and deleting cell specifications, which define the
+properties of battery cell types (manufacturer, chemistry, ratings, etc.).
+"""
+
 from typing import Any
 
 from pydantic import ValidationError
@@ -8,12 +15,42 @@ from .models import CellSpecification
 
 
 class CellSpecificationClient:
-    def __init__(self, client: Any):
+    """Client for managing cell specifications.
+
+    Provides methods to create, read, update, and delete cell specifications,
+    which define the properties of battery cell types (manufacturer, chemistry,
+    ratings, etc.).
+    """
+
+    def __init__(self, client: Any) -> None:
+        """Initialize the CellSpecificationClient.
+
+        Parameters
+        ----------
+        client : Any
+            The HTTP client instance for making API requests.
+        """
         self.client = client
 
     def get(self, cell_spec_id: str) -> CellSpecification:
-        """
-        Get a specific cell specification by ID.
+        """Get a specific cell specification by ID.
+
+        Parameters
+        ----------
+        cell_spec_id : str
+            The ID of the cell specification to retrieve.
+
+        Returns
+        -------
+        CellSpecification
+            The requested cell specification object.
+
+        Raises
+        ------
+        ValueError
+            If the response format is invalid.
+        RuntimeError
+            If the API call fails.
         """
         endpoint = f"/cell_specifications/{cell_spec_id}"
         try:
@@ -26,8 +63,19 @@ class CellSpecificationClient:
             raise RuntimeError(f"API call to {endpoint} failed: {e}") from e
 
     def list(self) -> list[CellSpecification]:
-        """
-        List all cell specs.
+        """List all cell specifications.
+
+        Returns
+        -------
+        list[CellSpecification]
+            A list of all cell specification objects.
+
+        Raises
+        ------
+        ValueError
+            If the response format is not a list or items are invalid.
+        RuntimeError
+            If the API call fails.
         """
         endpoint = "/cell_specifications"
         try:
@@ -45,17 +93,37 @@ class CellSpecificationClient:
             raise RuntimeError(f"API call to {endpoint} failed: {e}") from e
 
     def create(self, data: dict[str, Any]) -> CellSpecification:
-        """
-        Create a new cell spec.
+        """Create a new cell specification.
+
+        Parameters
+        ----------
+        data : dict[str, Any]
+            Dictionary containing the cell specification data.
+
+        Returns
+        -------
+        CellSpecification
+            The newly created cell specification object.
         """
         endpoint = "/cell_specifications"
         response_data = self.client.post(endpoint, data)
         return CellSpecification(**response_data)
 
     def create_or_get(self, data: dict[str, Any]) -> CellSpecification:
-        """
-        Create a new cell spec if it doesn't exist, otherwise get the
-        existing one.
+        """Create a new cell specification or get an existing one.
+
+        Creates a new cell specification if it doesn't exist, otherwise returns
+        the existing one.
+
+        Parameters
+        ----------
+        data : dict[str, Any]
+            Dictionary containing the cell specification data.
+
+        Returns
+        -------
+        CellSpecification
+            The cell specification object (newly created or existing).
         """
         try:
             return self.create(data)
@@ -64,39 +132,58 @@ class CellSpecificationClient:
                 existing_id = e.data.get("existing_cell_specification_id")
                 if existing_id:
                     return self.get(existing_id)
-            raise e
+            raise
 
     def update(self, cell_spec_id: str, data: dict[str, Any]) -> CellSpecification:
-        """
-        Update an existing cell specification.
+        """Update an existing cell specification.
 
         Parameters
         ----------
         cell_spec_id : str
             The ID of the cell specification to update.
-        data : dict
+        data : dict[str, Any]
             Dictionary containing the fields to update. Supports nested
             component/material data for upsert.
 
         Returns
         -------
         CellSpecification
-            The updated cell specification.
+            The updated cell specification object.
         """
         endpoint = f"/cell_specifications/{cell_spec_id}"
         response_data = self.client.put(endpoint, data)
         return CellSpecification(**response_data)
 
     def delete(self, cell_spec_id: str) -> None:
-        """
-        Delete a cell specification by ID.
+        """Delete a cell specification by ID.
+
+        Parameters
+        ----------
+        cell_spec_id : str
+            The ID of the cell specification to delete.
         """
         endpoint = f"/cell_specifications/{cell_spec_id}"
         self.client.delete(endpoint)
 
     def get_by_slug(self, cell_spec_slug: str) -> CellSpecification:
-        """
-        Get a specific cell specification by slug.
+        """Get a specific cell specification by slug.
+
+        Parameters
+        ----------
+        cell_spec_slug : str
+            The slug of the cell specification to retrieve.
+
+        Returns
+        -------
+        CellSpecification
+            The requested cell specification object.
+
+        Raises
+        ------
+        ValueError
+            If the response format is invalid.
+        RuntimeError
+            If the API call fails.
         """
         endpoint = f"/cell_specifications/slug/{cell_spec_slug}"
         try:

ionworks/client.py

@@ -1,3 +1,10 @@
+"""Main client module for the Ionworks API.
+
+This module provides the :class:`Ionworks` client, which is the main entry point
+for interacting with the Ionworks API. It handles authentication, request/response
+processing, and provides access to all API resources through sub-clients.
+"""
+
 import os
 from typing import Any, cast
 
@@ -11,21 +18,41 @@ from .errors import IonworksError
 from .job import JobClient
 from .pipeline import PipelineClient
 from .simulation import SimulationClient
+from .validators import set_dataframe_backend
 
 
 class Ionworks:
-    def __init__(self, api_key: str | None = None, api_url: str | None = None):
+    """Client for interacting with the Ionworks API.
+
+    Handles authentication, request/response processing, and provides access
+    to all API resources through sub-clients.
+    """
+
+    def __init__(
+        self,
+        api_key: str | None = None,
+        api_url: str | None = None,
+        dataframe_backend: str | None = None,
+    ) -> None:
         """Initialize Ionworks client.
 
         Parameters
         ----------
-        api_key : str, optional
-            API key. If not provided, will look for IONWORKS_API_KEY env var
-        api_url : str, optional
-            API URL. If not provided, will look for IONWORKS_API_URL env var
+        api_key : str | None
+            API key. If not provided, will look for IONWORKS_API_KEY env var.
+        api_url : str | None
+            API URL. If not provided, will look for IONWORKS_API_URL env var.
+        dataframe_backend : str | None
+            DataFrame backend for returned data: "polars" or "pandas".
+            If not provided, uses IONWORKS_DATAFRAME_BACKEND env var
+            (defaults to "polars").
         """
         load_dotenv()
 
+        # Set DataFrame backend (explicit param > env var > default "polars")
+        if dataframe_backend is not None:
+            set_dataframe_backend(dataframe_backend)
+
         self.api_key = api_key or os.getenv("IONWORKS_API_KEY")
         if not self.api_key:
             raise ValueError(
@@ -96,7 +123,7 @@ class Ionworks:
                 except requests.exceptions.JSONDecodeError:
                     error_msg += f": {e.response.text}"
             else:
-                error_msg += f": {str(e)}"
+                error_msg += f": {e!s}"
             # Raise the custom error for general request issues
             raise IonworksError(error_msg) from None
 
@@ -119,8 +146,10 @@ class Ionworks:
     def health_check(self) -> dict[str, Any]:
         """Check the health of the Ionworks API.
 
-        Returns:
-            dict[str, Any]: Health check response
+        Returns
+        -------
+        dict[str, Any]
+            Health check response.
         """
         response_data = self.get("/healthz")
         return cast(dict[str, Any], response_data)

ionworks/errors.py

@@ -1,3 +1,10 @@
+"""
+Custom exception classes for the Ionworks API client.
+
+This module defines :class:`IonworksError`, which is raised when API requests
+fail or return error responses.
+"""
+
 from typing import Any
 
 
@@ -7,18 +14,27 @@ class IonworksError(Exception):
     Attributes
     ----------
     message : str
-        A string description of the error
+        A string description of the error.
     data : dict[str, Any] | None
-        Structured error data if available (e.g., from API error response)
+        Structured error data if available (e.g., from API error response).
     status_code : int | None
-        HTTP status code if applicable
+        HTTP status code if applicable.
     """
 
     def __init__(
         self,
         message: str | dict[str, Any],
         status_code: int | None = None,
-    ):
+    ) -> None:
+        """Initialize the IonworksError.
+
+        Parameters
+        ----------
+        message : str | dict[str, Any]
+            Error message string or dict containing error details.
+        status_code : int | None
+            Optional HTTP status code.
+        """
         self.status_code = status_code
 
         # Parse message into string and optional data dict
@@ -31,5 +47,6 @@ class IonworksError(Exception):
 
         super().__init__(self.message)
 
-    def __str__(self):
+    def __str__(self) -> str:
+        """Return string representation of the error."""
         return f"error code: {self.status_code}, message: {self.message}"

ionworks/job.py

@@ -1,18 +1,26 @@
-from typing import Any, Optional
+"""Job client for managing asynchronous jobs.
+
+This module provides the :class:`JobClient` for submitting, monitoring, and
+managing background jobs in the Ionworks platform.
+"""
+
+from typing import Any
 
 from pydantic import BaseModel, ValidationError
 
 
-# Model for the data required to create ANY job
 class JobCreationPayload(BaseModel):
+    """Payload for creating a job."""
+
     job_type: str
     params: dict[str, Any]
     priority: int = 5  # Default priority
-    callback_url: Optional[str] = None  # Optional callback
+    callback_url: str | None = None  # Optional callback
 
 
-# Model for the detailed job response
 class JobResponse(BaseModel):
+    """Response model for job details."""
+
     job_id: str
     status: str
     job_type: str
@@ -20,18 +28,30 @@ class JobResponse(BaseModel):
     priority: int
     created_at: str
     updated_at: str
-    metadata: Optional[dict[str, Any]]
-    error: Optional[str]
-    result: Optional[dict[str, Any]]
+    metadata: dict[str, Any] | None
+    error: str | None
+    result: dict[str, Any] | None
 
 
 class JobClient:
-    def __init__(self, client: Any):
+    """Client for managing asynchronous jobs.
+
+    This class provides methods to create, retrieve, list, and cancel jobs
+    in the Ionworks platform.
+    """
+
+    def __init__(self, client: Any) -> None:
+        """Initialize the JobClient.
+
+        Parameters
+        ----------
+        client : Any
+            The HTTP client instance used for API requests.
+        """
         self.client = client
 
     def create(self, payload: JobCreationPayload) -> JobResponse:
-        """
-        Submit a job using the provided payload.
+        """Submit a job using the provided payload.
 
         Parameters
         ----------
@@ -59,53 +79,87 @@ class JobClient:
             return JobResponse(**response_data)
         except ValidationError as e:
             # Catch Pydantic validation errors specifically
-            raise ValueError(
-                f"Invalid response format received from {endpoint}: {e}"
-            ) from e
+            msg = f"Invalid response format received from {endpoint}: {e}"
+            raise ValueError(msg) from e
         # RequestExceptions (including HTTPError) are handled by client._post
 
     def get(self, job_id: str) -> JobResponse:
-        """
-        Get the status and details of a specific job.
+        """Get the status and details of a specific job.
+
+        Parameters
+        ----------
+        job_id : str
+            The ID of the job to retrieve.
+
+        Returns
+        -------
+        JobResponse
+            Job details and current status.
+
+        Raises
+        ------
+        ValueError
+            If the response parsing fails.
         """
         endpoint = f"/jobs/{job_id}"
         try:
             response_data = self.client.get(endpoint)
             return JobResponse(**response_data)
         except ValidationError as e:
-            raise ValueError(
-                f"Invalid response format received from {endpoint}: {e}"
-            ) from e
+            msg = f"Invalid response format received from {endpoint}: {e}"
+            raise ValueError(msg) from e
 
     def list(self) -> list[JobResponse]:
-        """
-        List all jobs.
+        """List all jobs.
+
+        Returns
+        -------
+        list[JobResponse]
+            List of all jobs with their details.
+
+        Raises
+        ------
+        ValueError
+            If the response is not a list or job data format is invalid.
         """
         endpoint = "/jobs/"
+        response_data = self.client.get(endpoint)
+        # Ensure response_data is a list before list comprehension
+        if not isinstance(response_data, list):
+            msg = (
+                f"Unexpected response format from {endpoint}: expected a list, "
+                f"got {type(response_data).__name__}"
+            )
+            raise ValueError(msg)
+        # Apply validation within list comprehension
         try:
-            response_data = self.client.get(endpoint)
-            # Ensure response_data is a list before list comprehension
-            if not isinstance(response_data, list):
-                raise ValueError(
-                    f"Unexpected response format from {endpoint}: expected a list, "
-                    "got {type(response_data).__name__}"
-                )
-            # Apply validation within list comprehension
             return [JobResponse(**job) for job in response_data]
         except ValidationError as e:
-            raise ValueError(
-                f"Invalid job data format received from {endpoint}: {e}"
-            ) from e
+            msg = f"Invalid job data format received from {endpoint}: {e}"
+            raise ValueError(msg) from e
 
     def cancel(self, job_id: str) -> JobResponse:
-        """
-        Cancel a job.
+        """Cancel a job.
+
+        Parameters
+        ----------
+        job_id : str
+            The ID of the job to cancel.
+
+        Returns
+        -------
+        JobResponse
+            Updated job details with cancelled status.
+
+        Raises
+        ------
+        ValueError
+            If the response parsing fails.
         """
         endpoint = f"/jobs/{job_id}/cancel"
         try:
             response_data = self.client.post(endpoint, json_payload={})
             return JobResponse(**response_data)
         except ValidationError as e:
-            raise ValueError(
-                f"Invalid response format received from {endpoint}: {e}"
-            ) from e
+            msg = f"Invalid response format received from {endpoint}: {e}"
+            raise ValueError(msg) from e

ionworks/models.py

@@ -1,5 +1,4 @@
-"""
-Pydantic models for the Ionworks API client.
+"""Pydantic models for the Ionworks API client.
 
 These models use extra="allow" to accept any fields from the API response,
 letting the API handle validation. Required fields are kept minimal.
@@ -7,17 +6,15 @@ letting the API handle validation. Required fields are kept minimal.
 
 from typing import Any
 
-import pandas as pd  # type: ignore
 from pydantic import BaseModel, ConfigDict, field_validator
 
-from .validators import dict_to_df_validator
+from .validators import DataFrame, dict_to_df_validator
 
 # --- Cell Specification Models --- #
 
 
 class CellSpecification(BaseModel):
-    """
-    Cell specification model - accepts any fields from the API.
+    """Cell specification model - accepts any fields from the API.
 
     The API returns nested component/material data and ratings objects.
     This model is permissive to allow the API to define the schema.
@@ -34,9 +31,7 @@ class CellSpecification(BaseModel):
 
 
 class CellInstance(BaseModel):
-    """
-    Cell instance model - accepts any fields from the API.
-    """
+    """Cell instance model - accepts any fields from the API."""
 
     model_config = ConfigDict(extra="allow")
 
@@ -48,9 +43,7 @@ class CellInstance(BaseModel):
 
 
 class CellMeasurement(BaseModel):
-    """
-    Cell measurement model - accepts any fields from the API.
-    """
+    """Cell measurement model - accepts any fields from the API."""
 
     model_config = ConfigDict(extra="allow")
 
@@ -94,18 +87,18 @@ class ConfirmUploadResponse(BaseModel):
 class CellMeasurementDetailBase(BaseModel):
     """Base model for measurement detail with steps and time series."""
 
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
     measurement: CellMeasurement
-    steps: pd.DataFrame
-    time_series: pd.DataFrame
+    steps: DataFrame
+    time_series: DataFrame
 
     @field_validator("steps", "time_series", mode="before")
     @classmethod
     def convert_dict_to_df(cls, v: Any) -> Any:
+        """Convert dictionary to DataFrame (polars or pandas based on config)."""
         return dict_to_df_validator(v)
 
-    class Config:
-        arbitrary_types_allowed = True  # Allow pandas DataFrame
-
 
 class CellInstanceDetail(BaseModel):
     """Detail model for a cell instance with all measurements."""
@@ -126,32 +119,32 @@ class CellMeasurementDetail(CellMeasurementDetailBase):
 
 
 class CellInstanceMeasurementsWithSteps(BaseModel):
-    """
-    Response model for all measurements + steps associated with a specific cell
-    instance.
-    Returns normalized data: separate lists for instance, measurements, and steps.
-    Includes parent information for consistency with other group endpoints.
+    """Response model for measurements + steps for a cell instance.
+
+    Returns normalized data: separate lists for instance, measurements,
+    and steps. Includes parent information for consistency with other
+    group endpoints.
     """
 
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
     specification: CellSpecification
     instance: CellInstance
     measurements: list[CellMeasurement]
-    steps: list[pd.DataFrame]
+    steps: list[DataFrame]
 
     @field_validator("steps", mode="before")
     @classmethod
     def convert_dict_to_df(cls, v: Any) -> Any:
+        """Convert dictionary or list of dictionaries to DataFrames."""
         if isinstance(v, list):
             return [dict_to_df_validator(item) for item in v]
         return dict_to_df_validator(v)
 
-    class Config:
-        arbitrary_types_allowed = True
-
 
 class CellSpecificationInstances(BaseModel):
-    """
-    Response model for all instances associated with a cell specification.
+    """Response model for all instances associated with a cell specification.
+
     Returns normalized data: separate lists for specification, instances,
     and measurements.
     """
@@ -162,24 +155,23 @@ class CellSpecificationInstances(BaseModel):
 
 
 class CellSpecificationInstancesWithSteps(BaseModel):
-    """
-    Response model for all instances + measurements + steps associated with a cell
-    specification.
-    Returns normalized data: separate lists for specification, instances, measurements,
-    and steps.
+    """Response model for instances + measurements + steps for a cell spec.
+
+    Returns normalized data: separate lists for specification, instances,
+    measurements, and steps.
     """
 
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
     specification: CellSpecification
     instances: list[CellInstance]
     measurements: list[CellMeasurement]
-    steps: list[pd.DataFrame]
+    steps: list[DataFrame]
 
     @field_validator("steps", mode="before")
     @classmethod
     def convert_dict_to_df(cls, v: Any) -> Any:
+        """Convert dictionary or list of dictionaries to DataFrames."""
         if isinstance(v, list):
             return [dict_to_df_validator(item) for item in v]
         return dict_to_df_validator(v)
-
-    class Config:
-        arbitrary_types_allowed = True