koleo-cli 0.2.137__tar.gz → 0.2.137.2__tar.gz

{koleo_cli-0.2.137 → koleo_cli-0.2.137.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: koleo-cli
-Version: 0.2.137
+Version: 0.2.137.2
 Summary: Koleo CLI
 Home-page: https://github.com/lzgirlcat/koleo-cli
 Author: Zoey !

{koleo_cli-0.2.137 → koleo_cli-0.2.137.2}/koleo/__init__.py

@@ -1,4 +1,2 @@
 from .api import KoleoAPI
-
-
-__all__ = ["KoleoAPI"]
+from .types import *

{koleo_cli-0.2.137 → koleo_cli-0.2.137.2}/koleo/api.py

@@ -1,53 +1,59 @@
 import typing as t
 from datetime import datetime
 
-from requests import Session, PreparedRequest, Response
+from requests import PreparedRequest, Response, Session
 
 from koleo.types import *
 
 
-class KoleoAPIException(Exception):
-    status: int
-    request: PreparedRequest
-    response: Response
 
-    def __init__(self, response: Response, *args: object) -> None:
-        super().__init__(*args)
-        self.status = response.status_code
-        self.request = response.request
-        self.response = response
 
-    @classmethod
-    def from_response(cls, response: Response) -> "KoleoAPIException":
-        if response.status_code == 404:
-            return KoleoNotFound(response)
-        elif response.status_code == 401:
-            return KoleoUnauthorized(response)
-        elif response.status_code == 403:
-            return KoleoForbidden(response)
-        elif response.status_code == 429:
-            return KoleoRatelimited(response)
-        else:
-            return KoleoAPIException(response)
 
 
-class KoleoNotFound(KoleoAPIException):
-    pass
 
+class errors:
+    class KoleoAPIException(Exception):
+        status: int
+        request: PreparedRequest
+        response: Response
 
-class KoleoForbidden(KoleoAPIException):
-    pass
+        def __init__(self, response: Response, *args: object) -> None:
+            super().__init__(*args)
+            self.status = response.status_code
+            self.request = response.request
+            self.response = response
 
+        @classmethod
+        def from_response(cls, response: Response) -> "t.Self":
+            if response.status_code == 404:
+                return errors.KoleoNotFound(response)
+            elif response.status_code == 401:
+                return errors.KoleoUnauthorized(response)
+            elif response.status_code == 403:
+                return errors.KoleoForbidden(response)
+            elif response.status_code == 429:
+                return errors.KoleoRatelimited(response)
+            else:
+                return cls(response)
 
-class KoleoUnauthorized(KoleoAPIException):
-    pass
+    class KoleoNotFound(KoleoAPIException):
+        pass
 
+    class KoleoForbidden(KoleoAPIException):
+        pass
 
-class KoleoRatelimited(KoleoAPIException):
-    pass
+
+    class KoleoUnauthorized(KoleoAPIException):
+        pass
+
+
+    class KoleoRatelimited(KoleoAPIException):
+        pass
 
 
 class KoleoAPI:
+    errors = errors
+
     def __init__(self) -> None:
         self.session = Session()
         self.base_url = "https://koleo.pl"
@@ -61,17 +67,23 @@ class KoleoAPI:
         headers = {**self.base_headers, **kwargs.get("headers", {})}
         r = self.session.get(self.base_url + path, *args, headers=headers, **kwargs)
         if not r.ok:
-            raise KoleoAPIException.from_response(r)
+            raise errors.KoleoAPIException.from_response(r)
         return r
 
     def _get_json(self, path, *args, **kwargs) -> t.Any:
         r = self._get(path, *args, **kwargs)
-        return r.json()
+        res = r.json()
+        if res is None:
+            raise self.errors.KoleoNotFound(r)
+        return res
 
     def _get_bytes(self, path, *args, **kwargs) -> bytes:
         r = self._get(path, *args, **kwargs)
         return r.content
 
+    def get_stations(self) -> list[ExtendedStationInfo]:
+        return self._get_json("/api/v2/main/stations")
+
     def find_station(self, query: str, language: str = "pl") -> list[ExtendedStationInfo]:
         # https://koleo.pl/ls?q=tere&language=pl
         return self._get_json("/ls", query={"q": query, "language": language})
@@ -106,11 +118,11 @@ class KoleoAPI:
         params = {"brand": brand_name, "nr": number}
         if name:
             params["name"] = name.upper()  # WHY!!!!!!!!!
-        return self._get_json("/api/v2/main/train_calendars", params=params)
+        return self._get_json("/pl/train_calendars", params=params)
 
     def get_train(self, id: int) -> TrainDetailResponse:
         # https://koleo.pl/pl/trains/142821312
-        return self._get_json(f"/api/v2/main/trains/{id}")
+        return self._get_json(f"/pl/trains/{id}")
 
     def get_connections(
         self,
@@ -119,7 +131,7 @@ class KoleoAPI:
         brand_ids: list[int],
         date: datetime,
         direct: bool = False,
-        purchasable: bool = False
+        purchasable: bool = False,
     ) -> ...:
         params = {
             "query[date]": date.strftime("%d-%m-%Y %H:%M:%S"),
@@ -127,7 +139,7 @@ class KoleoAPI:
             "query[end_station]": end,
             "query[only_purchasable]": purchasable,
             "query[direct]": direct,
-            "query[brand_ids]": brand_ids
+            "query[brand_ids][]": brand_ids,
         }
         return self._get_json("/api/v2/main/connections", params=params)
 

{koleo_cli-0.2.137 → koleo_cli-0.2.137.2}/koleo/cli.py

@@ -1,15 +1,16 @@
-from datetime import datetime
 from argparse import ArgumentParser
+from datetime import datetime, timedelta
 
 from rich.console import Console
 from rich.traceback import install
 
 from .api import KoleoAPI
-from .types import TrainOnStationInfo, TrainDetailResponse
-from .utils import name_to_slug, parse_datetime, time_dict_to_dt
-from .storage import Storage, DEFAULT_CONFIG_PATH
+from .storage import DEFAULT_CONFIG_PATH, Storage
+from .types import TrainDetailResponse, TrainOnStationInfo, ExtendedBaseStationInfo
+from .utils import convert_platform_number, name_to_slug, parse_datetime, arr_dep_to_dt
 
-install(show_locals=True)
+
+install(show_locals=False, max_frames=2)
 
 
 class CLI:
@@ -43,10 +44,15 @@ class CLI:
     def storage(self, storage: Storage):
         self._storage = storage
 
-    def list_stations(self, name: str): ...
+    def list_stations(self, name: str):
+        ...
 
     def get_departures(self, station_id: int, date: datetime):
-        trains = self.client.get_departures(station_id, date)
+        cache_id = f"dep-{station_id}-{date.strftime("%Y-%m-%d")}"
+        trains = (
+            self.storage.get_cache(cache_id) or
+            self.storage.set_cache(cache_id, self.client.get_departures(station_id, date))
+        )
         trains = [
             i
             for i in trains
@@ -57,7 +63,11 @@ class CLI:
         return table
 
     def get_arrivals(self, station_id: int, date: datetime):
-        trains = self.client.get_arrivals(station_id, date)
+        cache_id = f"arr-{station_id}-{date.strftime("%Y-%m-%d")}"
+        trains = (
+            self.storage.get_cache(cache_id) or
+            self.storage.set_cache(cache_id, self.client.get_arrivals(station_id, date))
+        )
         trains = [
             i
             for i in trains
@@ -68,43 +78,49 @@ class CLI:
         return table
 
     def full_departures(self, station: str, date: datetime):
-        slug = name_to_slug(station)
-        st = self.storage.get_cache(f"st-{slug}") or self.storage.set_cache(
-            f"st-{slug}", self.client.get_station_by_slug(slug)
-        )
+        st = self.get_station(station)
         station_info = f"[bold blue]{st["name"]}[/bold blue] ID: {st["id"]}"
         self.console.print(station_info)
         self.get_departures(st["id"], date)
 
     def full_arrivals(self, station: str, date: datetime):
-        slug = name_to_slug(station)
-        st = self.storage.get_cache(f"st-{slug}") or self.storage.set_cache(
-            f"st-{slug}", self.client.get_station_by_slug(slug)
-        )
+        st = self.get_station(station)
         station_info = f"[bold blue]{st["name"]}[/bold blue] ID: {st["id"]}"
         self.console.print(station_info)
         self.get_arrivals(st["id"], date)
 
     def train_info(self, brand: str, name: str, date: datetime):
         brand = brand.upper().strip()
-        name = name.upper().strip()
-        cache_id = f"tc-{brand}-{name}"
+        name = name.strip()
         if name.isnumeric():
             number = int(name)
-            name = ""
-        elif len(parts := name.split(" ")) == 2 or len(parts := name.split("-")) == 2:
-            number, name = parts
+            train_name = ""
+        elif len((parts := name.split(" "))) == 2 or len((parts := name.split("-"))) == 2:
+            number, train_name = parts
             number = int(number)
         else:
             raise ValueError("Invalid train name!")
+        brands = self.storage.get_cache("brands") or self.storage.set_cache("brands", self.client.get_brands())
+        if brand not in [i["name"] for i in brands]:
+            res = {i["logo_text"]: i["name"] for i in brands}.get(brand)
+            if not res:
+                raise ValueError("Invalid brand name!")
+            brand = res
+        cache_id = f"tc-{brand}-{number}-{name}"
         train_calendars = self.storage.get_cache(cache_id) or self.storage.set_cache(
-            cache_id, self.client.get_train_calendars(brand, number, name)
+            cache_id, self.client.get_train_calendars(brand, number, train_name)
         )
-        brands = self.storage.get_cache("brands") or self.storage.set_cache("brands", self.client.get_brands())
         train_id = train_calendars["train_calendars"][0]["date_train_map"][date.strftime("%Y-%m-%d")]
         train_details = self.client.get_train(train_id)
-        brand = next(iter(i for i in brands if i["id"] == train_details["train"]["brand_id"]), {}).get("name", "")
-        data = f"{brand} {train_details["train"]["train_full_name"]}\n"
+        brand = next(iter(i for i in brands if i["id"] == train_details["train"]["brand_id"]), {}).get("logo_text", "")
+        parts = [f"{brand} {train_details["train"]["train_full_name"]}"]
+        route_start = arr_dep_to_dt(train_details["stops"][0]["departure"])
+        route_end = arr_dep_to_dt(train_details["stops"][-1]["arrival"])
+        if route_end.hour < route_start.hour or (route_end.hour==route_start.hour and route_end.minute < route_end.minute):
+            route_end += timedelta(days=1)
+        travel_time = route_end - route_start
+        speed = train_details["stops"][-1]["distance"] / 1000 / travel_time.seconds * 3600
+        parts.append(f"[white]  {travel_time.seconds//3600}h{(travel_time.seconds % 3600)/60:.0f}m {speed:^4.1f}km/h [/white]")
         vehicle_types: dict[str, str] = {
             stop["station_display_name"]: stop["vehicle_type"]
             for stop in train_details["stops"]
@@ -115,20 +131,32 @@ class CLI:
             start = keys[0]
             for i in range(1, len(keys)):
                 if vehicle_types[keys[i]] != vehicle_types[start]:
-                    data += f"[bold green]  {start} - {keys[i]}:[/bold green] {vehicle_types[start]}\n"
+                    parts.append(f"[bold green]  {start} - {keys[i]}:[/bold green] {vehicle_types[start]}")
                     start = keys[i]
-            data += f"[bold green]  {start} - {keys[-1]}:[/bold green] {vehicle_types[start]}"
-        self.console.print(data)
+            parts.append(f"[bold green]  {start} - {keys[-1]}:[/bold green] {vehicle_types[start]}")
+        self.console.print("\n".join(parts))
         self.console.print(self.train_route_table(train_details))
 
-    def route(self, start: str, end: str, date: datetime, direct: bool = False, purchasable: bool = False):
-        slug = name_to_slug(station)
-        st = self.storage.get_cache(f"st-{slug}") or self.storage.set_cache(
-            f"st-{slug}", self.client.get_station_by_slug(slug)
+    def connections(self, start: str, end: str, date: datetime, brands: list[str], direct: bool = False, purchasable: bool = False):
+        start_station = self.get_station(start)
+        end_station = self.get_station(end)
+        brands = [i.lower().strip() for i in brands]
+        api_brands = self.storage.get_cache("brands") or self.storage.set_cache("brands", self.client.get_brands())
+        if not brands:
+            connection_brands = [i["id"] for i in api_brands]
+        else:
+            connection_brands = [i["id"] for i in api_brands if i["name"].lower().strip() in brands or i["logo_text"].lower().strip() in brands]
+            if not connection_brands:
+                self.console.print(f'[bold red]No brands match: "{', '.join(brands)}"[/bold red]')
+                exit(2)
+        connections = self.client.get_connections(
+            start_station["name_slug"],
+            end_station["name_slug"],
+            connection_brands,
+            date,
+            direct,
+            purchasable
         )
-        station_info = f"[bold blue]{st["name"]}[/bold blue] ID: {st["id"]}"
-        self.console.print(station_info)
-        self.get_arrivals(st["id"], date)
 
     def trains_on_station_table(self, trains: list[TrainOnStationInfo], type: int = 1):
         parts = []
@@ -136,22 +164,34 @@ class CLI:
         for train in trains:
             time = train["departure"] if type == 1 else train["arrival"]
             assert time
-            brand = next(iter(i for i in brands if i["id"] == train["brand_id"]), {}).get("name")
+            brand = next(iter(i for i in brands if i["id"] == train["brand_id"]), {}).get("logo_text")
+            platform = convert_platform_number(train["platform"]) if train["platform"] else ""
+            position_info = f"{platform}/{train["track"]}" if train["track"] else platform
             parts.append(
-                f"[bold green]{time[11:16]}[/bold green] {brand} {train["train_full_name"]}[purple] {train["stations"][0]["name"]} [/purple]"
+                f"[bold green]{time[11:16]}[/bold green] {brand} {train["train_full_name"]}[purple] {train["stations"][0]["name"]} {position_info}[/purple]"
             )
         return "\n".join(parts)
 
     def train_route_table(self, train: TrainDetailResponse):
         parts = []
         for stop in train["stops"]:
-            arr = time_dict_to_dt(stop["arrival"])
-            dep = time_dict_to_dt(stop["departure"])
+            arr = arr_dep_to_dt(stop["arrival"])
+            dep = arr_dep_to_dt(stop["departure"])
+            platform = convert_platform_number(stop["platform"]) or ""
             parts.append(
-                f"[white underline]{stop["distance"] / 1000:0.4}km[/white underline] [bold green]{arr.strftime("%H:%M")}[/bold green] - [bold red]{dep.strftime("%H:%M")}[/bold red] [purple] {stop["station_display_name"]} {stop["platform"]} [/purple]"
+                f"[white underline]{stop["distance"] / 1000:^5.1f}km[/white underline] [bold green]{arr.strftime("%H:%M")}[/bold green] - [bold red]{dep.strftime("%H:%M")}[/bold red] [purple]{stop["station_display_name"]} {platform} [/purple]"
             )
         return "\n".join(parts)
 
+    def get_station(self, station: str) -> ExtendedBaseStationInfo:
+        slug = name_to_slug(station)
+        try:
+            return self.storage.get_cache(f"st-{slug}") or self.storage.set_cache(
+                f"st-{slug}", self.client.get_station_by_slug(slug)
+            )
+        except self.client.errors.KoleoNotFound:
+            self.console.print(f'[bold red]Station not found: "{station}"[/bold red]')
+            exit(2)
 
 def main():
     cli = CLI()
@@ -159,9 +199,9 @@ def main():
     parser = ArgumentParser("koleo", description="Koleo CLI")
     parser.add_argument("-c", "--config", help="Custom config path.", default=DEFAULT_CONFIG_PATH)
     parser.add_argument("--nocolor", help="Disable color output", action="store_true", default=False)
-    subparsers = parser.add_subparsers(title="actions", required=False)
+    subparsers = parser.add_subparsers(title="actions", required=False) # type: ignore
 
-    departures = subparsers.add_parser("departures", aliases=["d", "dep"], help="Allows you to list station departures")
+    departures = subparsers.add_parser("departures", aliases=["d", "dep", "odjazdy", "o"], help="Allows you to list station departures")
     departures.add_argument(
         "station",
         help="The station name",
@@ -178,7 +218,7 @@ def main():
     departures.add_argument("-s", "--save", help="save the station as your default one", action="store_true")
     departures.set_defaults(func=cli.full_departures, pass_=["station", "date"])
 
-    arrivals = subparsers.add_parser("arrivals", aliases=["a", "arr"], help="Allows you to list station departures")
+    arrivals = subparsers.add_parser("arrivals", aliases=["a", "arr", "przyjazdy", "p"], help="Allows you to list station departures")
     arrivals.add_argument(
         "station",
         help="The station name",
@@ -197,7 +237,7 @@ def main():
 
     train_route = subparsers.add_parser(
         "trainroute",
-        aliases=["r", "tr", "t"],
+        aliases=["r", "tr", "t", "poc", "pociąg"],
         help="Allows you to show the train's route",
     )
     train_route.add_argument("brand", help="The brand name", type=str)
@@ -211,6 +251,40 @@ def main():
     )
     train_route.set_defaults(func=cli.train_info, pass_=["brand", "name", "date"])
 
+    connections = subparsers.add_parser(
+        "connections",
+        aliases=["do", "z", "szukaj", "path", "find"],
+        help="Allows you to search for connections from a to b",
+    )
+    connections.add_argument("start", help="The starting station", type=str)
+    connections.add_argument("end", help="The end station", type=str)
+    connections.add_argument(
+        "-d",
+        "--date",
+        help="the date",
+        type=lambda s: parse_datetime(s),
+        default=datetime.now(),
+    )
+    connections.add_argument(
+        "-b",
+        "--brands",
+        help="Brands to include",
+        action="extend", nargs="+", type=str, default=[]
+    )
+    connections.add_argument(
+        "-f",
+        "--direct",
+        help="whether or not the result should only include direct trains",
+        action="store_true", default=False
+    )
+    connections.add_argument(
+        "-p",
+        "--purchasable",
+        help="whether or not the result should only trains purchasable on koleo",
+        action="store_true", default=False
+    )
+    connections.set_defaults(func=cli.connections, pass_=["start", "end", "brands", "date", "direct" "purchasable"])
+
     args = parser.parse_args()
 
     storage = Storage.load(path=args.config)

{koleo_cli-0.2.137 → koleo_cli-0.2.137.2}/koleo/storage.py

@@ -1,9 +1,10 @@
 import typing as t
-from time import time
-from dataclasses import dataclass, asdict, field
-from sys import platform
+from dataclasses import asdict, dataclass, field
 from json import dump, load
-from os import makedirs, path as ospath
+from os import makedirs
+from os import path as ospath
+from sys import platform
+from time import time
 
 
 def get_adequate_config_path():
@@ -42,7 +43,7 @@ class Storage:
         else:
             data = {}
         storage = cls(**data)
-        storage._path = path
+        storage._path = expanded
         return storage
 
     def get_cache(self, id: str) -> t.Any | None:
@@ -62,10 +63,9 @@ class Storage:
         return item
 
     def save(self):
-        expanded = ospath.expanduser(self._path)
-        dir = ospath.dirname(expanded)
+        dir = ospath.dirname(self._path)
         if dir:
             if not ospath.exists(dir):
                 makedirs(dir)
-        with open(expanded, "w+") as f:
-            dump(asdict(self), f)
+        with open(self._path, "w+") as f:
+            dump(asdict(self), f, indent=True)

{koleo_cli-0.2.137 → koleo_cli-0.2.137.2}/koleo/types.py

@@ -12,6 +12,7 @@ class ExtendedBaseStationInfo(BaseStationInfo):
     hits: int
     version: str  # "A", "B"
     is_group: bool
+    city: str | None
     region: str
     country: str
     latitude: float
@@ -175,3 +176,9 @@ class TrainStop(t.TypedDict):
 class TrainDetailResponse(t.TypedDict):
     train: TrainDetail
     stops: list[TrainStop]
+
+
+class ConnectiontrainDetail(TrainDetail):
+
+    arrival: TimeDict | str # WTF KOLEO!!!!
+    departure: TimeDict | str # WTF KOLEO!!!!

koleo_cli-0.2.137.2/koleo/utils.py

@@ -0,0 +1,65 @@
+from datetime import datetime, time, timedelta
+
+from .types import TimeDict
+
+
+def parse_datetime(s: str):
+    try:
+        now = datetime.now()
+        dt = datetime.strptime(s, "%d-%m")
+        return dt.replace(year=now.year, hour=0, minute=0)
+    except ValueError:
+        pass
+    try:
+        return datetime.strptime(s, "%Y-%m-%d").replace(hour=0, minute=0)
+    except ValueError:
+        pass
+    if s[0] == "+":
+        return datetime.now().replace(hour=0, minute=0) + timedelta(days=int(s[1:]))
+    return datetime.combine(datetime.now(), datetime.strptime(s, "%H:%M").time())
+
+
+def arr_dep_to_dt(i: TimeDict | str):
+    if isinstance(i, str):
+        return datetime.fromisoformat(i)
+    now = datetime.today()
+    return datetime.combine(now, time(i["hour"], i["minute"], i["second"]))
+
+
+TRANSLITERATIONS = {
+    "ł": "l",
+    "ń": "n",
+    "ą": "a",
+    "ę": "e",
+    "ś": "s",
+    "ć": "c",
+    "ó": "o",
+    "ź": "z",
+    "ż": "z",
+    " ": "-",
+    "_": "-",
+}
+
+
+def name_to_slug(name: str) -> str:
+    return "".join([TRANSLITERATIONS.get(char, char) for char in name.lower()])
+
+
+NUMERAL_TO_ARABIC = {
+    "I": 1,
+    "II": 2,
+    "III": 3,
+    "IV": 4,
+    "V": 5,
+    "VI": 6,
+    "VII": 7,
+    "VIII": 8,
+    "IX": 9,
+    "X": 10,
+    "XI": 11,  # wtf poznań???
+    "XII": 12,  # just to be safe
+}
+
+
+def convert_platform_number(number: str) -> int | None:
+    return NUMERAL_TO_ARABIC.get(number)

{koleo_cli-0.2.137 → koleo_cli-0.2.137.2}/koleo_cli.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: koleo-cli
-Version: 0.2.137
+Version: 0.2.137.2
 Summary: Koleo CLI
 Home-page: https://github.com/lzgirlcat/koleo-cli
 Author: Zoey !

{koleo_cli-0.2.137 → koleo_cli-0.2.137.2}/setup.py

@@ -14,7 +14,7 @@ def parse_requirements_file(path):
 
 setuptools.setup(
     name="koleo-cli",
-    version="0.2.137",
+    version="0.2.137.2",
     description="Koleo CLI",
     long_description=long_description(),
     long_description_content_type="text/markdown",

koleo_cli-0.2.137/koleo/utils.py

@@ -1,41 +0,0 @@
-from datetime import datetime, time
-
-from .types import TimeDict
-
-
-def parse_datetime(s: str):
-    now = datetime.today()
-    try:
-        dt = datetime.strptime(s, "%d-%m")
-        return dt.replace(year=now.year, hour=0, minute=0)
-    except ValueError:
-        pass
-    try:
-        return datetime.strptime(s, "%Y-%m-%d").replace(hour=0, minute=0)
-    except ValueError:
-        pass
-    return datetime.combine(now, datetime.strptime(s, "%H:%M").time())
-
-
-def time_dict_to_dt(s: TimeDict):
-    now = datetime.today()
-    return datetime.combine(now, time(s["hour"], s["minute"], s["second"]))
-
-
-TRANSLITERATIONS = {
-    "ł": "l",
-    "ń": "n",
-    "ą": "a",
-    "ę": "e",
-    "ś": "s",
-    "ć": "c",
-    "ó": "o",
-    "ź": "z",
-    "ż": "z",
-    " ": "-",
-    "_": "-",
-}
-
-
-def name_to_slug(name: str) -> str:
-    return "".join([TRANSLITERATIONS.get(char, char) for char in name.lower()])