File size: 1,447 Bytes
d6ea71e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 |
"""Schema for SPADL actions."""
from typing import Any, Optional
import pandera as pa
from pandera.typing import Series
from . import config as spadlconfig
class SPADLSchema(pa.SchemaModel):
"""Definition of a 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()
start_x: Series[float] = pa.Field(ge=0, le=spadlconfig.field_length)
start_y: Series[float] = pa.Field(ge=0, le=spadlconfig.field_width)
end_x: Series[float] = pa.Field(ge=0, le=spadlconfig.field_length)
end_y: Series[float] = pa.Field(ge=0, 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)
result_id: Series[int] = pa.Field(isin=spadlconfig.results_df().result_id)
result_name: Optional[Series[str]] = pa.Field(isin=spadlconfig.results_df().result_name)
class Config: # noqa: D106
strict = True
coerce = True
|