Skip to content

Contracts

Contracts define the shared column and metadata objects used by dataset configs, processing stages, and generated outputs.

MappingSpec

Bases: str

Canonical column name with dtype, aliases, parser, and optional values.

MappingSpec is a str subclass, so it can be passed directly to Polars selectors and aliases while carrying schema metadata for validation and raw column discovery. Common columns live in BaseColumns; dataset mapping modules derive dataset-specific variants with aliases or parsers.

Examples:

>>> curr = MappingSpec("Current [A]", dtype=pl.Float64)
>>> str(curr)
'Current [A]'
>>> curr.dtype
Float64
>>> raw_time = BaseColumns.time.with_alias("Test_Time(s)", "Test Time (s)")
>>> raw_time.matching_name(["Test Time (s)"])
'Test Time (s)'
>>> parsed_time = raw_time.with_parser(lambda col: pl.col(col).cast(pl.Float64))
>>> parsed_time.parser is not None
True
>>> git_status = BaseColumns.git_status
>>> git_status.values.clean
'clean'

Methods:

  • with_alias

    Return a copy with dataset-specific source column aliases.

  • with_parser

    Return a copy with a raw-column parser expression factory.

  • with_values

    Return a copy with fixed named values attached.

  • matching_name

    Return the first available column matching this spec or its aliases.

Attributes:

  • values (ValuesT) –

    Fixed named values declared for this column.

values property

values: ValuesT

Fixed named values declared for this column.

Returns:

  • ValuesT

    The values namespace attached to this spec.

Raises:

  • TypeError

    If no values namespace was declared.

with_alias

with_alias(*alias: str) -> MappingSpec[ValuesT]

Return a copy with dataset-specific source column aliases.

Parameters:

  • alias (str, default: () ) –

    Raw/source column names that should match this canonical spec.

Returns:

  • MappingSpec[ValuesT]

    A spec with the same canonical name and metadata, but the supplied

  • MappingSpec[ValuesT]

    source aliases.

Source code in batgrad/contracts/mapping.py
def with_alias(self, *alias: str) -> MappingSpec[ValuesT]:
    """Return a copy with dataset-specific source column aliases.

    Args:
        alias: Raw/source column names that should match this canonical spec.

    Returns:
        A spec with the same canonical name and metadata, but the supplied
        source aliases.
    """
    return MappingSpec[ValuesT](
        str(self),
        dtype=self.dtype,
        alias=alias,
        description=self.description,
        parser=self.parser,
        values=self._values,
    )

with_parser

with_parser(
    parser: Callable[[str], Expr],
) -> MappingSpec[ValuesT]

Return a copy with a raw-column parser expression factory.

Parameters:

  • parser (Callable[[str], Expr]) –

    Callable receiving the matched raw column name and returning a Polars expression for the parsed canonical column.

Returns:

  • MappingSpec[ValuesT]

    A spec with the same canonical name, aliases, and values, but the

  • MappingSpec[ValuesT]

    supplied parser.

Source code in batgrad/contracts/mapping.py
def with_parser(self, parser: Callable[[str], pl.Expr]) -> MappingSpec[ValuesT]:
    """Return a copy with a raw-column parser expression factory.

    Args:
        parser: Callable receiving the matched raw column name and returning
            a Polars expression for the parsed canonical column.

    Returns:
        A spec with the same canonical name, aliases, and values, but the
        supplied parser.
    """
    return MappingSpec[ValuesT](
        str(self),
        dtype=self.dtype,
        alias=self.alias[1:],
        description=self.description,
        parser=parser,
        values=self._values,
    )

with_values

with_values(values: NewValuesT) -> MappingSpec[NewValuesT]

Return a copy with fixed named values attached.

Parameters:

  • values (NewValuesT) –

    Namespace of allowed or conventional values for the column.

Returns:

  • MappingSpec[NewValuesT]

    A spec with the same canonical name, aliases, parser, and dtype, but

  • MappingSpec[NewValuesT]

    the supplied values namespace.

Source code in batgrad/contracts/mapping.py
def with_values[NewValuesT](self, values: NewValuesT) -> MappingSpec[NewValuesT]:
    """Return a copy with fixed named values attached.

    Args:
        values: Namespace of allowed or conventional values for the column.

    Returns:
        A spec with the same canonical name, aliases, parser, and dtype, but
        the supplied values namespace.
    """
    return MappingSpec[NewValuesT](
        str(self),
        dtype=self.dtype,
        alias=self.alias[1:],
        description=self.description,
        parser=self.parser,
        values=values,
    )

matching_name

matching_name(
    columns: set[str] | tuple[str, ...] | list[str],
) -> str | None

Return the first available column matching this spec or its aliases.

Parameters:

  • columns (set[str] | tuple[str, ...] | list[str]) –

    Available column names to search. Matching is case-insensitive and checks the canonical name first, followed by aliases in declaration order.

Returns:

  • str | None

    The original available column name when a match is found, otherwise

  • str | None

    None.

Source code in batgrad/contracts/mapping.py
def matching_name(self, columns: set[str] | tuple[str, ...] | list[str]) -> str | None:
    """Return the first available column matching this spec or its aliases.

    Args:
        columns: Available column names to search. Matching is
            case-insensitive and checks the canonical name first, followed by
            aliases in declaration order.

    Returns:
        The original available column name when a match is found, otherwise
        `None`.
    """
    available = {column.casefold(): column for column in columns}
    for alias in self.alias:
        match = available.get(alias.casefold())
        if match is not None:
            return match
    return None

MetadataLayout dataclass

MetadataLayout(
    required: tuple[MappingSpec, ...] = (),
    optional: Mapping[MappingSpec, object | None] = dict(),
)

Required columns and optional default values for stage metadata.

Layouts are used for both manifest rows and parquet footers. required columns must be present; optional entries provide default metadata values that can be added or overridden by dataset/protocol-specific configuration.

Attributes:

  • required (tuple[MappingSpec, ...]) –

    Metadata columns that must be present.

  • optional (Mapping[MappingSpec, object | None]) –

    Metadata columns and default values that may be added by the stage, dataset, or protocol configuration.

Examples:

>>> layout = MetadataLayout(required=(BaseColumns.proto,))
>>> expanded = layout.with_optional({BaseColumns.cell_id: None})
>>> expanded.columns
('protocol', 'cell id')
>>> expanded.values
{'cell id': None}

Methods:

  • with_optional

    Return a layout with additional or overridden optional defaults.

columns property

columns: tuple[MappingSpec, ...]

Required columns followed by optional columns.

Returns:

values property

values: dict[MappingSpec, object | None]

Optional metadata defaults as a mutable copy.

Returns:

with_optional

with_optional(
    optional: Mapping[MappingSpec, object | None],
) -> MetadataLayout

Return a layout with additional or overridden optional defaults.

Parameters:

  • optional (Mapping[MappingSpec, object | None]) –

    Optional metadata columns and default values to merge into this layout.

Returns:

  • MetadataLayout

    A layout with the same required columns and merged optional defaults.

Source code in batgrad/contracts/metadata.py
def with_optional(self, optional: Mapping[MappingSpec, object | None]) -> MetadataLayout:
    """Return a layout with additional or overridden optional defaults.

    Args:
        optional: Optional metadata columns and default values to merge into
            this layout.

    Returns:
        A layout with the same required columns and merged optional defaults.
    """
    values = dict(self.optional)
    values.update(optional)
    return MetadataLayout(
        required=self.required,
        optional=values,
    )

StageLayout dataclass

StageLayout(
    stage_id: DatasetStageId,
    manifest: MetadataLayout,
    footer: MetadataLayout,
)

Manifest and footer metadata contract for one processing stage.

Manifest metadata is validated against and written to manifest rows. Footer metadata is encoded as parquet key/value metadata when data files are written. Optional values set to None declare metadata that is expected from task metadata rather than a fixed stage-level default.

Attributes:

  • stage_id (DatasetStageId) –

    Pipeline stage this layout describes.

  • manifest (MetadataLayout) –

    Metadata required or defaulted in the stage manifest rows.

  • footer (MetadataLayout) –

    Metadata required or defaulted in generated parquet footers.

Examples:

>>> layout = INGEST_STAGE_METADATA.with_manifest(
...     {BaseColumns.cell_id: None, BaseColumns.cidx: None}
... ).with_footer({BaseColumns.time_conv: "start_of_interval"})
>>> BaseColumns.cell_id in layout.manifest.optional
True
>>> layout.footer.values[BaseColumns.time_conv]
'start_of_interval'

Methods:

  • with_manifest

    Return a stage layout with extra optional manifest metadata.

  • with_footer

    Return a stage layout with extra optional footer metadata.

with_manifest

with_manifest(
    optional: Mapping[MappingSpec, object | None],
) -> StageLayout

Return a stage layout with extra optional manifest metadata.

Parameters:

Returns:

  • StageLayout

    A stage layout with updated manifest metadata and unchanged footer

  • StageLayout

    metadata.

Source code in batgrad/contracts/metadata.py
def with_manifest(self, optional: Mapping[MappingSpec, object | None]) -> StageLayout:
    """Return a stage layout with extra optional manifest metadata.

    Args:
        optional: Manifest metadata columns and default values to add or
            override.

    Returns:
        A stage layout with updated manifest metadata and unchanged footer
        metadata.
    """
    return StageLayout(
        stage_id=self.stage_id,
        manifest=self.manifest.with_optional(optional),
        footer=self.footer,
    )
with_footer(
    optional: Mapping[MappingSpec, object | None],
) -> StageLayout

Return a stage layout with extra optional footer metadata.

Parameters:

Returns:

  • StageLayout

    A stage layout with unchanged manifest metadata and updated footer

  • StageLayout

    metadata.

Source code in batgrad/contracts/metadata.py
def with_footer(self, optional: Mapping[MappingSpec, object | None]) -> StageLayout:
    """Return a stage layout with extra optional footer metadata.

    Args:
        optional: Footer metadata columns and default values to add or
            override.

    Returns:
        A stage layout with unchanged manifest metadata and updated footer
        metadata.
    """
    return StageLayout(
        stage_id=self.stage_id,
        manifest=self.manifest,
        footer=self.footer.with_optional(optional),
    )