atomic/spadl/schema.py

"""Schema for Atomic-SPADL actions."""

from typing import Any, Optional

import pandera as pa
from pandera.typing import Series

from . import config as spadlconfig


class AtomicSPADLSchema(pa.SchemaModel):
    """Definition of an Atomic-SPADL dataframe."""

    game_id: Series[Any] = pa.Field()
    original_event_id: Series[Any] = pa.Field(nullable=True)
    action_id: Series[int] = pa.Field()
    period_id: Series[int] = pa.Field(ge=1, le=5)
    time_seconds: Series[float] = pa.Field(ge=0)
    team_id: Series[Any] = pa.Field()
    player_id: Series[Any] = pa.Field()
    x: Series[float] = pa.Field(ge=0, le=spadlconfig.field_length)
    y: Series[float] = pa.Field(ge=0, le=spadlconfig.field_width)
    dx: Series[float] = pa.Field(ge=-spadlconfig.field_length, le=spadlconfig.field_length)
    dy: Series[float] = pa.Field(ge=-spadlconfig.field_width, le=spadlconfig.field_width)
    bodypart_id: Series[int] = pa.Field(isin=spadlconfig.bodyparts_df().bodypart_id)
    bodypart_name: Optional[Series[str]] = pa.Field(isin=spadlconfig.bodyparts_df().bodypart_name)
    type_id: Series[int] = pa.Field(isin=spadlconfig.actiontypes_df().type_id)
    type_name: Optional[Series[str]] = pa.Field(isin=spadlconfig.actiontypes_df().type_name)

    class Config:  # noqa: D106
        strict = True
        coerce = True