Source code for ionworks.models

"""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.
"""

from typing import Any

from pydantic import BaseModel, ConfigDict, field_validator

from .validators import DataFrame, dict_to_df_validator

# --- Cell Specification Models --- #




class CellSpecification(BaseModel):
    """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.
    """

    model_config = ConfigDict(extra="allow")

    id: str
    name: str



# --- Cell Instance Models --- #




class CellInstance(BaseModel):
    """Cell instance model - accepts any fields from the API."""

    model_config = ConfigDict(extra="allow")

    id: str
    name: str
    cell_specification_id: str



# --- Cell Measurement Models --- #




class CellMeasurement(BaseModel):
    """Cell measurement model - accepts any fields from the API."""

    model_config = ConfigDict(extra="allow")

    id: str
    name: str
    cell_instance_id: str



# --- Bundle Models --- #




class CellMeasurementBundleResponse(BaseModel):
    """Response from creating a measurement bundle."""

    measurement: CellMeasurement
    steps_created: int
    file_path: str





class InitiateUploadResponse(BaseModel):
    """Response from the initiate-upload endpoint for signed URL uploads."""

    measurement_id: str
    signed_url: str
    token: str
    path: str





class ConfirmUploadResponse(BaseModel):
    """Response from the confirm-upload endpoint after successful upload."""

    instance: CellInstance
    measurement: CellMeasurement
    steps_created: int
    file_path: str



# --- Detail Models --- #




class CellMeasurementDetail(BaseModel):
    """Detail model for a measurement with steps and time series.

    Returns minimal data by default: foreign keys for parent objects
    rather than nested objects. Use the spec/instance clients to
    fetch parent objects if needed.
    """

    model_config = ConfigDict(arbitrary_types_allowed=True)

    measurement: CellMeasurement
    specification_id: str | None = None
    instance_id: str | None = None
    steps: DataFrame | None = None
    time_series: DataFrame | None = None
    cycles: DataFrame | None = None



    @field_validator("steps", "time_series", "cycles", 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 CellInstanceDetail(BaseModel):
    """Detail model for a cell instance with all measurements.

    Returns a foreign key for the parent specification rather
    than a nested object. Use ``client.cell_spec.get(detail
    .specification_id)`` to fetch the full specification.
    """

    instance: CellInstance
    specification_id: str
    measurements: list[CellMeasurementDetail]





class StepsAndCycles(BaseModel):
    """Steps and cycle metrics for a measurement.

    Returned by the ``/steps_and_cycles`` endpoint which
    fetches both in one call (cycles are derived from steps).
    """

    model_config = ConfigDict(arbitrary_types_allowed=True)

    steps: DataFrame
    cycles: DataFrame



    @field_validator("steps", "cycles", mode="before")
    @classmethod
    def convert_dict_to_df(cls, v: Any) -> Any:
        """Convert dictionary to DataFrame."""
        return dict_to_df_validator(v)