koleo-cli 0.2.137.17__py3-none-any.whl → 0.2.137.18__py3-none-any.whl

koleo/__init__.py

@@ -1,2 +1,2 @@
 from .api import KoleoAPI
-from .types import *
+from .api.types import *

koleo/api/__init__.py

@@ -0,0 +1,2 @@
+from .client import KoleoAPI
+from .types import *

koleo/api/base.py

@@ -0,0 +1,70 @@
+from asyncio import sleep as asleep
+
+from aiohttp import (
+    ClientConnectorError,
+    ClientOSError,
+    ClientResponse,
+    ClientResponseError,
+    ClientSession,
+)
+from orjson import loads
+
+from .logging import LoggingMixin
+
+
+class JsonableData(bytes):
+    response: ClientResponse
+
+    def __new__(cls, *args, response: ClientResponse, **kwargs):
+        obj = super().__new__(cls, *args, **kwargs)
+        obj.response = response
+        return obj
+
+    def json(self):
+        if "_json" not in self.__dict__:
+            self._json = loads(bytes(self))
+        return self._json
+
+
+class BaseAPIClient(LoggingMixin):
+    _session: ClientSession
+
+    exc = ClientResponseError
+
+    @property
+    def session(self) -> "ClientSession":
+        if not hasattr(self, "_session"):
+            self._session = ClientSession()
+        return self._session
+
+    async def close(self):
+        return await self.session.close()
+
+    async def exc_getter(self, r: ClientResponse) -> Exception | None:
+        return
+
+    async def request(self, method, url: str, *args, retries: int = 4, fail_wait: float = 8, **kwargs) -> JsonableData:
+        try:
+            async with self.session.request(method, url, *args, **kwargs) as r:
+                if not r.ok:
+                    self.dl(r.headers)
+                    try:
+                        self.dl(await r.text())
+                    except UnicodeDecodeError:
+                        self.dl("Response is not text!")
+                    if exc := (await self.exc_getter(r)):
+                        raise exc
+                    r.raise_for_status()
+                return JsonableData(await r.read(), response=r)
+        except (ClientConnectorError, ClientOSError) as e:
+            if retries > 0:
+                await asleep(fail_wait)
+                return await self.request(
+                    method,
+                    url,
+                    *args,
+                    retries=retries - 1,
+                    fail_wait=fail_wait,
+                    **kwargs,
+                )
+            raise e

koleo/api/client.py

@@ -0,0 +1,221 @@
+import typing as t
+from datetime import datetime
+
+from aiohttp import ClientResponse
+
+from koleo.api.types import *
+
+from .base import BaseAPIClient
+from .errors import errors
+
+
+class KoleoAPI(BaseAPIClient):
+    errors = errors
+
+    def __init__(self, auth: dict[str, str] | None = None) -> None:
+        self.base_url = "https://koleo.pl"
+        self.version = 2
+        self.base_headers = {
+            "x-koleo-version": str(self.version),
+            "User-Agent": "Koleo-CLI(https://pypi.org/project/koleo-cli)",
+        }
+        self._auth: dict[str, str] | None = auth
+        self._auth_valid: bool | None = None
+
+    async def get(self, path, use_auth: bool = False, *args, **kwargs):
+        headers = {**self.base_headers, **kwargs.pop("headers", {})}
+        if self._auth and use_auth:
+            headers["cookie"] = "; ".join([f"{k}={v}" for k, v in self._auth.items()])
+        r = await self.request("GET", self.base_url + path, headers=headers, *args, **kwargs)
+        if len(r) == 0:
+            raise self.errors.KoleoNotFound(r.response)
+        return r
+
+    async def post(self, path, use_auth: bool = False, *args, **kwargs):
+        headers = {**self.base_headers, **kwargs.pop("headers", {})}
+        if self._auth and use_auth:
+            headers["cookie"] = ("; ".join([f"{k}={v}" for k, v in self._auth.items()]),)
+        r = await self.request("POST", self.base_url + path, headers=headers, *args, **kwargs)
+        if len(r) == 0:
+            raise self.errors.KoleoNotFound(r.response)
+        return r
+
+    async def exc_getter(self, r: ClientResponse) -> Exception | None:
+        return await self.errors.from_response(r)
+
+    async def _require_auth(self) -> t.Literal[True]:
+        if self._auth is None:
+            raise errors.AuthRequired()
+        if self._auth_valid is None:
+            await self.get_current_session()
+            self._auth_valid = True
+        return True
+
+    async def get_stations(self) -> list[ExtendedStationInfo]:
+        return (await self.get("/api/v2/main/stations")).json()
+
+    async def find_station(self, query: str, language: str = "pl") -> list[SearchStationInfo]:
+        # https://koleo.pl/ls?q=tere&language=pl
+        return (await self.get("/ls", params={"q": query, "language": language})).json()["stations"]
+
+    async def get_station_by_id(self, id: int) -> ExtendedStationInfo:
+        # https://koleo.pl/api/v2/main/stations/by_id/24000
+        return (
+            await self.get(
+                f"/api/v2/main/stations/by_id/{id}",
+            )
+        ).json()
+
+    async def get_station_by_slug(self, slug: str) -> ExtendedStationInfo:
+        # https://koleo.pl/api/v2/main/stations/by_slug/inowroclaw
+        return (
+            await self.get(
+                f"/api/v2/main/stations/by_slug/{slug}",
+            )
+        ).json()
+
+    async def get_station_info_by_slug(self, slug: str) -> StationDetails:
+        # https://koleo.pl/api/v2/main/station_info/inowroclaw
+        return (
+            await self.get(
+                f"/api/v2/main/station_info/{slug}",
+            )
+        ).json()
+
+    async def get_departures(self, station_id: int, date: datetime) -> list[TrainOnStationInfo]:
+        # https://koleo.pl/api/v2/main/timetables/18705/2024-03-25/departures
+        return (
+            await self.get(
+                f"/api/v2/main/timetables/{station_id}/{date.strftime("%Y-%m-%d")}/departures",
+            )
+        ).json()
+
+    async def get_arrivals(self, station_id: int, date: datetime) -> list[TrainOnStationInfo]:
+        # https://koleo.pl/api/v2/main/timetables/18705/2024-03-25/arrivals
+        return (
+            await self.get(
+                f"/api/v2/main/timetables/{station_id}/{date.strftime("%Y-%m-%d")}/arrivals",
+            )
+        ).json()
+
+    async def get_train_calendars(self, brand_name: str, number: int, name: str | None = None) -> TrainCalendarResponse:
+        # https://koleo.pl/pl/train_calendars?brand=REG&nr=10417
+        # https://koleo.pl/pl/train_calendars?brand=IC&nr=1106&name=ESPERANTO ; WHY!!!! WHY!!!!!!1
+        params = {"brand": brand_name, "nr": number}
+        if name:
+            params["name"] = name.upper()  # WHY!!!!!!!!!
+        return (await self.get("/pl/train_calendars", params=params)).json()
+
+    async def get_train(self, id: int) -> TrainDetailResponse:
+        # https://koleo.pl/pl/trains/142821312
+        return (await self.get(f"/pl/trains/{id}")).json()
+
+    async def get_connections(
+        self,
+        start: str,
+        end: str,
+        brand_ids: list[int],
+        date: datetime,
+        direct: bool = False,
+        purchasable: bool = False,
+    ) -> list[ConnectionDetail]:
+        params = {
+            "query[date]": date.strftime("%d-%m-%Y %H:%M:%S"),
+            "query[start_station]": start,
+            "query[end_station]": end,
+            "query[only_purchasable]": str(purchasable).lower(),
+            "query[only_direct]": str(direct).lower(),
+            "query[brand_ids][]": brand_ids,
+        }
+        return (await self.get("/api/v2/main/connections", params=params)).json()["connections"]
+
+    async def get_connection(self, id: int) -> ConnectionDetail:
+        return (
+            await self.get(
+                f"/api/v2/main/connections/{id}",
+            )
+        ).json()
+
+    async def get_brands(self) -> list[ApiBrand]:
+        # https://koleo.pl/api/v2/main/brands
+        return (
+            await self.get(
+                "/api/v2/main/brands",
+            )
+        ).json()
+
+    async def get_carriers(self) -> list[Carrier]:
+        # https://koleo.pl/api/v2/main/carriers
+        return (
+            await self.get(
+                "/api/v2/main/carriers",
+            )
+        ).json()
+
+    async def get_discounts(self) -> list[DiscountInfo]:
+        # https://koleo.pl/api/v2/main/discounts
+        return (
+            await self.get(
+                "/api/v2/main/discounts",
+            )
+        ).json()
+
+    async def get_nested_train_place_types(self, connection_id: int) -> SeatsAvailabilityResponse:
+        # https://koleo.pl/api/v2/main/seats_availability/connection_id/train_nr/place_type
+        await self._require_auth()
+        if self._auth and "_koleo_token" not in self._auth:
+            res = await self.post(f"/prices/{connection_id}/passengers")
+            self._auth["_koleo_token"] = koleo_token = res.response.cookies["_koleo_token"].value
+        return (
+            await self.get(
+                f"/api/v2/main/nested_train_place_types/{connection_id}",
+                headers={"Authorization": f"Bearer {koleo_token}"},
+                use_auth=True,
+            )
+        ).json()
+
+    async def get_seats_availability(
+        self, connection_id: int, train_nr: int, place_type: int
+    ) -> SeatsAvailabilityResponse:
+        # https://koleo.pl/api/v2/main/seats_availability/connection_id/train_nr/place_type
+        return (
+            await self.get(
+                f"/api/v2/main/seats_availability/{connection_id}/{train_nr}/{place_type}",
+            )
+        ).json()
+
+    async def get_train_composition(
+        self, connection_id: int, train_nr: int, place_type: int
+    ) -> SeatsAvailabilityResponse:
+        # https://koleo.pl/api/v2/main/train_composition/connection_id/train_nr/place_type
+        return (
+            await self.get(
+                f"/api/v2/main/train_composition/{connection_id}/{train_nr}/{place_type}",
+            )
+        ).json()
+
+    async def get_carriage_type(self, id: int) -> CarriageType:
+        # https://koleo.pl/api/v2/main/carriage_types/id
+        return (
+            await self.get(
+                f"/api/v2/main/carriage_types/{id}",
+            )
+        ).json()
+
+    async def get_carriage_types(self) -> list[CarriageType]:
+        return (await self.get("/api/v2/main/carriage_types")).json()
+
+    async def get_station_keywoards(self) -> list[StationKeyword]:
+        return (await self.get("/api/v2/main/station_keywords")).json()
+
+    async def get_price(self, connection_id: int) -> Price | None:
+        res = await self.get(
+            f"/pl/prices/{connection_id}",
+        )
+        return res.json().get("price")
+
+    async def get_current_session(self) -> CurrentSession:
+        return (await self.get(f"/sessions/current", use_auth=True)).json()
+
+    async def get_current_user(self) -> CurrentUser:
+        return (await self.get(f"/users/current", use_auth=True)).json()

koleo/api/errors.py

@@ -0,0 +1,46 @@
+from typing import TYPE_CHECKING
+
+
+if TYPE_CHECKING:
+    from aiohttp import ClientResponse, RequestInfo
+
+
+class errors:
+    class KoleoAPIException(Exception):
+        status: int
+        request: "RequestInfo"
+        response: "ClientResponse"
+
+        def __init__(self, response: "ClientResponse", *args: object) -> None:
+            super().__init__(*args)
+            self.status = response.status
+            self.request = response.request_info
+            self.response = response
+
+    @staticmethod
+    async def from_response(response: "ClientResponse") -> "KoleoAPIException":
+        if response.status == 404:
+            return errors.KoleoNotFound(response)
+        elif response.status == 401:
+            return errors.KoleoUnauthorized(response)
+        elif response.status == 403:
+            return errors.KoleoForbidden(response)
+        elif response.status == 429:
+            return errors.KoleoRatelimited(response)
+        else:
+            return errors.KoleoAPIException(response, await response.text())
+
+    class KoleoNotFound(KoleoAPIException):
+        pass
+
+    class KoleoForbidden(KoleoAPIException):
+        pass
+
+    class KoleoUnauthorized(KoleoAPIException):
+        pass
+
+    class KoleoRatelimited(KoleoAPIException):
+        pass
+
+    class AuthRequired(Exception):
+        pass

koleo/api/logging.py

@@ -0,0 +1,56 @@
+import logging
+
+
+class LoggingMixin:
+    _l: logging.Logger
+    _l_name: str | None = None
+
+    def __init__(self, name: str | None = None) -> None:
+        if name:
+            self._l_name = name
+
+    @property
+    def logger(self) -> logging.Logger:
+        if not getattr(self, "_l", None):
+            self._l = logging.getLogger(self.logger_name)
+        return self._l
+
+    @property
+    def logger_name(self) -> str:
+        return getattr(self, "_l_name", None) or self.__class__.__name__.lower()
+
+    def dl(self, msg, *args, **kwargs):
+        self.logger.debug(msg, *args, **kwargs)
+
+    def error(self, msg, *args, **kwargs):
+        self.logger.error(msg, *args, **kwargs)
+
+    def warn(self, msg, *args, **kwargs):
+        self.logger.warning(msg, *args, **kwargs)
+
+    def info(self, msg, *args, **kwargs):
+        self.logger.info(msg, *args, **kwargs)
+
+    def create_logging_context(self, prefix: str) -> "ContextLogger":
+        return ContextLogger(self.logger, prefix)
+
+
+class ContextLogger:
+    def __init__(self, logger: logging.Logger, prefix: str) -> None:
+        self.logger = logger
+        self.prefix = prefix
+
+    def _make_msg(self, msg: str):
+        return f"{self.prefix}: {msg}"
+
+    def dl(self, msg, *args, **kwargs):
+        self.logger.debug(self._make_msg(msg), *args, **kwargs)
+
+    def info(self, msg, *args, **kwargs):
+        self.logger.info(self._make_msg(msg), *args, **kwargs)
+
+    def error(self, msg, *args, **kwargs):
+        self.logger.error(self._make_msg(msg), *args, **kwargs)
+
+    def warn(self, msg, *args, **kwargs):
+        self.logger.warning(self._make_msg(msg), *args, **kwargs)