hass-sncf/sncf/objects.py

import enum
from pydantic import BaseModel, validator, Field
import datetime
from typing import List, Optional

BaseModel.Config.extra = "forbid"


class Coord(BaseModel):
    lat: float
    lon: float


class Code(BaseModel):
    type: str
    value: str

class CO2Emission(BaseModel):
    value: float
    unit: str


class Link(BaseModel):
    internal: Optional[bool]
    type: str
    id: Optional[str]
    rel: Optional[str]
    templated: Optional[bool]
    href: Optional[str]


class AdministrativeRegion(BaseModel):
    id: str
    name: str
    label: str
    coord: Coord
    level: int
    zip_code: str
    insee: int


class TravelMode(BaseModel):
    id: str
    name: str

    @property
    def value(self) -> str:
        return self.id.split(":")[-1]


class PhysicalMode(TravelMode):
    @validator("id")
    def validate_physicalmode(cls, v):
        for physicalmode in PhysicalModes:
            if v == f"physical_mode:{physicalmode}":
                return v
        raise ValueError(v)


class PhysicalModes(str, enum.Enum):
    Air = "Air"
    Boat = "Boat"
    Bus = "Bus"
    BusRapidTransit = "BusRapidTransit"
    Coach = "Coach"
    Ferry = "Ferry"
    Funicular = "Funicular"
    LocalTrain = "LocalTrain"
    LongDistanceTrain = "LongDistanceTrain"
    Metro = "Metro"
    RailShuttle = "RailShuttle"
    RapidTransit = "RapidTransit"
    Shuttle = "Shuttle"
    SuspendedCableCar = "SuspendedCableCar"
    Taxi = "Taxi"
    Train = "Train"
    Tramway = "Tramway"


class StopArea(BaseModel):
    id: str
    name: str
    label: str
    coord: Coord
    administrative_regions: Optional[List[AdministrativeRegion]]
    stop_points: Optional[List["StopPoint"]]
    codes: List[Code]
    links: List[Link]
    timezone: str
    commercial_modes: Optional[List[TravelMode]]
    physical_modes: Optional[List[PhysicalMode]]


class StopPoint(BaseModel):
    id: str
    name: str
    coord: Coord
    administrative_regions: Optional[List[AdministrativeRegion]]
    equipments: List[str]
    stop_area: Optional[StopArea]
    links: List[Link]
    label: str


class JourneyStatus(str, enum.Enum):
    NONE = ""
    NO_SERVICE = "NO_SERVICE"
    REDUCED_SERVICE = "REDUCED_SERVICE"
    SIGNIFICANT_DELAYS = "SIGNIFICANT_DELAYS"
    DETOUR = "DETOUR"
    ADDITIONAL_SERVICE = "ADDITIONAL_SERVICE"
    MODIFIED_SERVICE = "MODIFIED_SERVICE"
    OTHER_EFFECT = "OTHER_EFFECT"
    UNKNOWN_EFFECT = "UNKNOWN_EFFECT"
    STOP_MOVED = "STOP_MOVED"


class SectionType(str, enum.Enum):
    public_transport = "public_transport"
    street_network = "street_network"
    waiting = "waiting"
    stay_in = "stay_in"
    transfer = "transfer"
    crow_fly = "crow_fly"
    on_demand_transport = "on_demand_transport"
    bss_rent = "bss_rent"
    bss_put_back = "bss_put_back"
    boarding = "boarding"
    landing = "landing"
    alighting = "alighting"
    park = "park"
    ridesharing = "ridesharing"


class SectionMode(str, enum.Enum):
    NONE = ""
    Walking = "walking"
    Bike = "bike"
    Car = "car"
    Taxi = "taxi"


class EmbeddedType(str, enum.Enum):
    administrative_region = "administrative_region"
    stop_area = "stop_area"
    stop_point = "stop_point"
    address = "address"
    poi = "poi"


class POIType(BaseModel):
    id: str
    name: str


class POI(BaseModel):
    id: str
    name: str
    label: str
    type: POIType
    # stands: todo


class Address(BaseModel):
    id: str
    name: str
    label: str
    coord: Coord
    house_number: int
    administrative_regions: List[AdministrativeRegion]


class Place(BaseModel):
    id: str
    name: str
    quality: int
    embedded_type: EmbeddedType
    administrative_region: Optional[AdministrativeRegion]
    stop_area: Optional[StopArea]
    stop_point: Optional[StopPoint]
    address: Optional[Address]
    poi: Optional[POI]

    @property
    def embed(self) -> AdministrativeRegion | StopArea | POI | Address | StopPoint:
        return getattr(self, self.embedded_type)


class DisplayInformation(BaseModel):
    network: str
    physical_mode: str
    commercial_mode: str
    code: str
    color: str
    text_color: str
    direction: str
    headsign: str
    label: str
    name: str
    trip_short_name: str
    equipments: List[str]
    description: str
    links: List[Link]


class Path(BaseModel):
    length: int
    name: str
    duration: int
    direction: int


class TransferType(str, enum.Enum):
    walking = "walking"
    stay_in = "stay_in"


class StopDateTime(BaseModel):
    stop_point: StopPoint
    additional_informations: List[str]
    base_departure_date_time: Optional[datetime.datetime]
    departure_date_time: Optional[datetime.datetime]
    base_arrival_date_time: Optional[datetime.datetime]
    arrival_date_time: Optional[datetime.datetime]
    links: List[str]

    @validator(
        "base_departure_date_time",
        "departure_date_time",
        "base_arrival_date_time",
        "arrival_date_time",
        pre=True, allow_reuse=True)
    def convert_datetime(cls, v):
        return datetime.datetime.strptime(v, "%Y%m%dT%H%M%S")


class DataFreshness(str, enum.Enum):
    realtime = "realtime"
    base_schedule = "base_schedule"


class Section(BaseModel):
    type: SectionType
    id: str
    mode: Optional[SectionMode]
    duration: int
    from_place: Optional[Place] = Field(alias='from')
    to_place: Optional[Place] = Field(alias='to')
    links: List[Link]
    display_informations: Optional[DisplayInformation]
    additional_informations: Optional[List[str]]  # todo: enum
    geojson: Optional[dict]  # todo
    path: Optional[List[Path]]
    transfer_type: Optional[TransferType]
    stop_date_times: Optional[List[StopDateTime]]
    departure_date_time: datetime.datetime
    arrival_date_time: datetime.datetime
    co2_emission: CO2Emission
    data_freshness: Optional[DataFreshness]
    base_departure_date_time: Optional[datetime.datetime]
    base_arrival_date_time: Optional[datetime.datetime]

    @validator(
        "departure_date_time",
        "arrival_date_time",
        "base_departure_date_time",
        "base_arrival_date_time",
        pre=True, allow_reuse=True)
    def convert_datetime(cls, v):
        return datetime.datetime.strptime(v, "%Y%m%dT%H%M%S")


class Cost(BaseModel):
    value: float
    currency: Optional[str]


class Fare(BaseModel):
    total: Cost
    found: bool
    links: List[Link]


class Durations(BaseModel):
    taxi: int
    walking: int
    car: int
    ridesharing: int
    bike: int
    total: Optional[int]


class Period(BaseModel):
    begin: datetime.date
    end: datetime.date
    @validator(
        "begin",
        "end",
        pre=True, allow_reuse=True)
    def convert_datetime(cls, v):
        return datetime.datetime.strptime(v, "%Y%m%d")


class WeekPattern(BaseModel):
    monday: bool
    tuesday: bool
    wednesday: bool
    thursday: bool
    friday: bool
    saturday: bool
    sunday: bool


class Calendar(BaseModel):
    active_periods: List[Period]
    week_pattern: WeekPattern

class Journey(BaseModel):
    duration: int
    nb_transfers: int
    departure_date_time: datetime.datetime
    requested_date_time: datetime.datetime
    arrival_date_time: datetime.datetime
    sections: List[Section]
    links: List[Link]
    type: str  # todo enum
    fare: Fare
    tags: List[str]
    status: JourneyStatus
    from_place: Place = None
    to: Place = None
    durations: Durations
    distances: Durations
    co2_emission: CO2Emission
    calendars: List[Calendar]

    @validator(
        "departure_date_time",
        "requested_date_time",
        "arrival_date_time",
        pre=True, allow_reuse=True)
    def convert_datetime(cls, v):
        return datetime.datetime.strptime(v, "%Y%m%dT%H%M%S")