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

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

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

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

@@ -6,11 +6,6 @@ from requests import PreparedRequest, Response, Session
 from koleo.types import *
 
 
-
-
-
-
-
 class errors:
     class KoleoAPIException(Exception):
         status: int
@@ -86,7 +81,13 @@ class KoleoAPI:
 
     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})
+        return self._get_json("/ls", params={"q": query, "language": language})["stations"]
+
+    def get_station_by_id(self, id: int) -> ExtendedBaseStationInfo:
+        # https://koleo.pl/api/v2/main/stations/by_id/24000
+        return self._get_json(
+            f"/api/v2/main/stations/by_id/{id}",
+        )
 
     def get_station_by_slug(self, slug: str) -> ExtendedBaseStationInfo:
         # https://koleo.pl/api/v2/main/stations/by_slug/inowroclaw
@@ -132,16 +133,16 @@ class KoleoAPI:
         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]": purchasable,
-            "query[direct]": direct,
+            "query[only_purchasable]": str(direct).lower(),
+            "query[only_direct]": str(direct).lower(),
             "query[brand_ids][]": brand_ids,
         }
-        return self._get_json("/api/v2/main/connections", params=params)
+        return self._get_json("/api/v2/main/connections", params=params)["connections"]
 
     def get_brands(self) -> list[ApiBrand]:
         # https://koleo.pl/api/v2/main/brands

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

@@ -1,3 +1,4 @@
+import re
 from argparse import ArgumentParser
 from datetime import datetime, timedelta
 
@@ -22,7 +23,15 @@ class CLI:
     ) -> None:
         self._client = client
         self._storage = storage
-        self.console = Console(color_system="standard", no_color=no_color)
+        self.no_color = no_color
+        self.console = Console(color_system="standard", no_color=no_color, highlight=False)
+
+    def print(self, text, *args, **kwargs):
+        if self.no_color:
+            result = re.sub(r'\[[^\]]*\]', '', text)
+            print(result)
+        else:
+            self.console.print(text, *args, **kwargs)
 
     @property
     def client(self) -> KoleoAPI:
@@ -44,9 +53,6 @@ class CLI:
     def storage(self, storage: Storage):
         self._storage = storage
 
-    def list_stations(self, name: str):
-        ...
-
     def get_departures(self, station_id: int, date: datetime):
         cache_id = f"dep-{station_id}-{date.strftime("%Y-%m-%d")}"
         trains = (
@@ -59,7 +65,7 @@ class CLI:
             if datetime.fromisoformat(i["departure"]).timestamp() > date.timestamp()  # type: ignore
         ]
         table = self.trains_on_station_table(trains)
-        self.console.print(table)
+        self.print(table)
         return table
 
     def get_arrivals(self, station_id: int, date: datetime):
@@ -74,21 +80,32 @@ class CLI:
             if datetime.fromisoformat(i["arrival"]).timestamp() > date.timestamp()  # type: ignore
         ]
         table = self.trains_on_station_table(trains, type=2)
-        self.console.print(table)
+        self.print(table)
         return table
 
     def full_departures(self, station: str, date: datetime):
         st = self.get_station(station)
-        station_info = f"[bold blue]{st["name"]}[/bold blue] ID: {st["id"]}"
-        self.console.print(station_info)
+        station_info = f"[bold blue][link=https://koleo.pl/dworzec-pkp/{st["name_slug"]}/{date.strftime("%Y-%m-%d")}]{st["name"]} at {date.strftime("%d-%m %H:%M")}[/bold blue] ID: {st["id"]}[/link]"
+        self.print(station_info)
         self.get_departures(st["id"], date)
 
     def full_arrivals(self, station: str, date: datetime):
         st = self.get_station(station)
-        station_info = f"[bold blue]{st["name"]}[/bold blue] ID: {st["id"]}"
-        self.console.print(station_info)
+        station_info = f"[bold blue][link=https://koleo.pl/dworzec-pkp/{st["name_slug"]}/{date.strftime("%Y-%m-%d")}]{st["name"]} at {date.strftime("%d-%m %H:%M")}[/bold blue] ID: {st["id"]}[/link]"
+        self.print(station_info)
         self.get_arrivals(st["id"], date)
 
+    def find_station(self, query: str | None):
+        if query:
+            stations = self.client.find_station(query)
+        else:
+            stations = (
+                self.storage.get_cache("stations") or
+                self.storage.set_cache("stations", self.client.get_stations())
+            )
+        for st in stations:
+            self.print(f"[bold blue][link=https://koleo.pl/dworzec-pkp/{st["name_slug"]}]{st["name"]}[/bold blue] ID: {st["id"]}[/link]")
+
     def train_info(self, brand: str, name: str, date: datetime):
         brand = brand.upper().strip()
         name = name.strip()
@@ -113,7 +130,7 @@ class CLI:
         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("logo_text", "")
-        parts = [f"{brand} {train_details["train"]["train_full_name"]}"]
+        parts = [f"[red]{brand}[/red] [bold blue]{train_details["train"]["train_full_name"]}[/bold blue]"]
         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):
@@ -131,13 +148,13 @@ class CLI:
             start = keys[0]
             for i in range(1, len(keys)):
                 if vehicle_types[keys[i]] != vehicle_types[start]:
-                    parts.append(f"[bold green]  {start} - {keys[i]}:[/bold green] {vehicle_types[start]}")
+                    parts.append(f"  {start} - {keys[i]}: [bold green]{vehicle_types[start]}[/bold green]")
                     start = keys[i]
-            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))
+            parts.append(f"  {start} - {keys[-1]}: [bold green]{vehicle_types[start]}[/bold green]")
+        self.print("\n".join(parts))
+        self.print(self.train_route_table(train_details))
 
-    def connections(self, start: str, end: str, date: datetime, brands: list[str], direct: bool = False, purchasable: bool = False):
+    def connections(self, start: str, end: str, date: datetime, brands: list[str], direct: bool, purchasable: bool):
         start_station = self.get_station(start)
         end_station = self.get_station(end)
         brands = [i.lower().strip() for i in brands]
@@ -147,7 +164,7 @@ class CLI:
         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]')
+                self.print(f'[bold red]No brands match: "{', '.join(brands)}"[/bold red]')
                 exit(2)
         connections = self.client.get_connections(
             start_station["name_slug"],
@@ -157,6 +174,59 @@ class CLI:
             direct,
             purchasable
         )
+        parts = [f"[bold blue]{start_station["name"]} → {end_station["name"]} at {date.strftime("%H:%M %d-%m")}[/bold blue]"]
+        for i in connections:
+            arr = arr_dep_to_dt(i["arrival"])
+            dep = arr_dep_to_dt(i["departure"])
+            travel_time = (arr - dep).seconds
+            parts.append(f"[bold green][link=https://koleo.pl/travel-options/{i["id"]}]{dep.strftime("%H:%M")} - {arr.strftime("%H:%M")}[/bold green] {travel_time//3600}h{(travel_time % 3600)/60:.0f}m {i['distance']}km: [/link]")
+            if i["constriction_info"]:
+                for constriction in i["constriction_info"]:
+                    parts.append(f" [bold red]- {constriction} [/bold red]")
+            if len(i["trains"]) == 1:
+                train = i["trains"][0]
+                stop = next(iter(i for i in train["stops"] if i["station_id"] == train["start_station_id"]), {})
+                stop_station = (
+                    start_station if stop["station_id"] == start_station["id"] else
+                    self.storage.get_cache(f"st-{stop['station_id']}")
+                    or self.storage.set_cache(f"st-{stop['station_id']}", self.client.get_station_by_id(stop['station_id']))
+                )
+                platform = convert_platform_number(stop["platform"]) if stop["platform"] else ""
+                position_info = f"{platform}/{stop["track"]}" if stop["track"] else platform
+                brand = next(iter(i for i in api_brands if i["id"] == train["brand_id"]), {}).get("logo_text")
+                parts[-1]+=f" [red]{brand}[/red] {train["train_full_name"]}[purple] {stop_station['name']} {position_info}[/purple]"
+            else:
+                previous_arrival: datetime | None = None
+                for train in i["trains"]:
+                    brand = next(iter(i for i in api_brands if i["id"] == train["brand_id"]), {}).get("logo_text")
+
+                    fs = next(iter(i for i in train["stops"] if i["station_id"] == train["start_station_id"]), {})
+                    fs_station = (
+                        start_station if fs["station_id"] == start_station["id"] else
+                        self.storage.get_cache(f"st-{fs['station_id']}")
+                        or self.storage.set_cache(f"st-{fs['station_id']}", self.client.get_station_by_id(fs['station_id']))
+                    )
+                    fs_platform = convert_platform_number(fs["platform"]) if fs["platform"] else ""
+                    fs_arr = arr_dep_to_dt(fs["arrival"])
+                    fs_dep = arr_dep_to_dt(fs["departure"])
+                    fs_info = f"[bold green]{fs_dep.strftime("%H:%M")} [/bold green][purple]{fs_station['name']} {f"{fs_platform}/{fs["track"]}" if fs["track"] else fs_platform}[/purple]"
+
+                    ls = next(iter(i for i in train["stops"] if i["station_id"] == train["end_station_id"]), {})
+                    ls_station = (
+                        start_station if ls["station_id"] == start_station["id"] else
+                        self.storage.get_cache(f"st-{ls['station_id']}")
+                        or self.storage.set_cache(f"st-{ls['station_id']}", self.client.get_station_by_id(ls['station_id']))
+                    )
+                    ls_platform = convert_platform_number(ls["platform"]) if ls["platform"] else ""
+                    ls_arr = arr_dep_to_dt(ls["arrival"])
+                    ls_dep = arr_dep_to_dt(ls["departure"])
+                    ls_info = f"[bold green]{ls_arr.strftime("%H:%M")} [/bold green][purple]{ls_station['name']} {f"{ls_platform}/{ls["track"]}" if ls["track"] else ls_platform}[/purple]"
+                    connection_time = (fs_dep - previous_arrival).seconds if previous_arrival else ""
+                    previous_arrival = ls_arr
+                    if connection_time:
+                        parts.append(f"  {connection_time//3600}h{(connection_time % 3600)/60:.0f}m at [purple]{fs_station['name']}[/purple]")
+                    parts.append(f"  [red]{brand}[/red] {train["train_full_name"]} {fs_info} - {ls_info}")
+        self.print("\n".join(parts))
 
     def trains_on_station_table(self, trains: list[TrainOnStationInfo], type: int = 1):
         parts = []
@@ -168,7 +238,7 @@ class CLI:
             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"]} {position_info}[/purple]"
+                f"[bold green]{time[11:16]}[/bold green] [red]{brand}[/red] {train["train_full_name"]}[purple] {train["stations"][0]["name"]} {position_info}[/purple]"
             )
         return "\n".join(parts)
 
@@ -190,7 +260,7 @@ class CLI:
                 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]')
+            self.print(f'[bold red]Station not found: "{station}"[/bold red]')
             exit(2)
 
 def main():
@@ -251,9 +321,18 @@ def main():
     )
     train_route.set_defaults(func=cli.train_info, pass_=["brand", "name", "date"])
 
+    stations = subparsers.add_parser("stations", aliases=["s", "find", "f", "stacje", "ls"], help="Allows you to find stations by their name")
+    stations.add_argument(
+        "query",
+        help="The station name",
+        default=None,
+        nargs="?",
+    )
+    stations.set_defaults(func=cli.find_station, pass_=["query"])
+
     connections = subparsers.add_parser(
         "connections",
-        aliases=["do", "z", "szukaj", "path", "find"],
+        aliases=["do", "z", "szukaj", "path"],
         help="Allows you to search for connections from a to b",
     )
     connections.add_argument("start", help="The starting station", type=str)
@@ -283,7 +362,7 @@ def main():
         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"])
+    connections.set_defaults(func=cli.connections, pass_=["start", "end", "brands", "date", "direct", "purchasable"])
 
     args = parser.parse_args()
 
@@ -291,6 +370,7 @@ def main():
     client = KoleoAPI()
     cli.client, cli.storage = client, storage
     cli.console.no_color = args.nocolor
+    cli.no_color = args.nocolor
     if hasattr(args, "station") and args.station is None:
         args.station = storage.favourite_station
     elif hasattr(args, "station") and args.save:

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

@@ -178,7 +178,60 @@ class TrainDetailResponse(t.TypedDict):
     stops: list[TrainStop]
 
 
+class ConnectionTrainStop(t.TypedDict):
+    arrival: TimeDict | str # WTF KOLEO!!!!
+    departure: TimeDict | str # WTF KOLEO!!!!
+    distance: int
+    in_path: bool
+    station_id: int
+    next_day: bool
+    position: int
+    train_nr: int
+    brand_id: int
+    entry_only: bool
+    exit_only: bool
+    platform: str
+    track: str
+    on_demand: bool
+
+
 class ConnectiontrainDetail(TrainDetail):
+    arrival: TimeDict | str # WTF KOLEO!!!!
+    departure: TimeDict | str # WTF KOLEO!!!!
+    stops: list[ConnectionTrainStop]
+    bookable: bool
+    train_attribute_ids: list[int]
+    travel_time: int
+    direction: str
+    start_station_id: int
+    end_station_id: int
+    fixed_carriage_composition: bool
+    is_option_groups_available: bool
 
+
+class ErrorDict(t.TypedDict):
+    type: str
+    value: str
+
+
+class ConnectionDetail(t.TypedDict):
+    id: int
+    distance: int
+    purchasable: bool
+    purchasable_errors: list[ErrorDict]
+    travel_time: int
+    changes: int
+    needs_document: bool
+    brand_ids: list[int]
+    start_station_id: int
+    end_station_id: int
     arrival: TimeDict | str # WTF KOLEO!!!!
     departure: TimeDict | str # WTF KOLEO!!!!
+    bookable: bool
+    special_event_slug: str | None
+    is_advanced_travel_options: bool
+    is_child_birthday_required: bool
+    max_passengers_count: bool
+    constriction_info: list[str]
+    is_estimated_timetable_available: bool
+    trains: list[ConnectiontrainDetail]

{koleo_cli-0.2.137.2 → koleo_cli-0.2.137.4}/koleo/utils.py

@@ -16,6 +16,12 @@ def parse_datetime(s: str):
         pass
     if s[0] == "+":
         return datetime.now().replace(hour=0, minute=0) + timedelta(days=int(s[1:]))
+    try:
+        now = datetime.now()
+        dt = datetime.strptime(s, "%d-%m %H:%M")
+        return dt.replace(year=now.year)
+    except ValueError:
+        pass
     return datetime.combine(datetime.now(), datetime.strptime(s, "%H:%M").time())
 
 
@@ -58,8 +64,9 @@ NUMERAL_TO_ARABIC = {
     "X": 10,
     "XI": 11,  # wtf poznań???
     "XII": 12,  # just to be safe
+    "BUS": "BUS"
 }
 
 
-def convert_platform_number(number: str) -> int | None:
+def convert_platform_number(number: str) -> int | None | str:
     return NUMERAL_TO_ARABIC.get(number)

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

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

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

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