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

koleo/cli.py

@@ -1,550 +0,0 @@
-import re
-from argparse import ArgumentParser
-from datetime import datetime, timedelta
-
-from rich.console import Console
-
-from .api import KoleoAPI
-from .storage import DEFAULT_CONFIG_PATH, Storage
-from .types import ExtendedBaseStationInfo, TrainDetailResponse, TrainOnStationInfo, TrainCalendar
-from .utils import RemainderString, arr_dep_to_dt, convert_platform_number, name_to_slug, parse_datetime
-
-
-class CLI:
-    def __init__(
-        self,
-        no_color: bool = False,
-        client: KoleoAPI | None = None,
-        storage: Storage | None = None,
-    ) -> None:
-        self._client = client
-        self._storage = storage
-        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:
-        if not self._client:
-            raise ValueError("Client not set!")
-        return self._client
-
-    @client.setter
-    def client(self, client: KoleoAPI):
-        self._client = client
-
-    @property
-    def storage(self) -> Storage:
-        if not self._storage:
-            raise ValueError("Storage not set!")
-        return self._storage
-
-    @storage.setter
-    def storage(self, storage: Storage):
-        self._storage = storage
-
-    def get_departures(self, station_id: int, date: datetime):
-        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
-            if datetime.fromisoformat(i["departure"]).timestamp() > date.timestamp()  # type: ignore
-        ]
-        table = self.trains_on_station_table(trains)
-        self.print(table)
-        return table
-
-    def get_arrivals(self, station_id: int, date: datetime):
-        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
-            if datetime.fromisoformat(i["arrival"]).timestamp() > date.timestamp()  # type: ignore
-        ]
-        table = self.trains_on_station_table(trains, type=2)
-        self.print(table)
-        return table
-
-    def full_departures(self, station: str, date: datetime):
-        st = self.get_station(station)
-        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][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 all_trains(self, station: str, date: datetime):
-        st = self.get_station(station)
-        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)
-        arr_cache_id = f"arr-{st['id']}-{date.strftime("%Y-%m-%d")}"
-        dep_cache_id = f"dep-{st['id']}-{date.strftime("%Y-%m-%d")}"
-        arrivals = self.storage.get_cache(arr_cache_id) or self.storage.set_cache(
-            arr_cache_id, self.client.get_arrivals(st["id"], date)
-        )
-        departures = self.storage.get_cache(dep_cache_id) or self.storage.set_cache(
-            dep_cache_id, self.client.get_departures(st["id"], date)
-        )
-        trains = sorted(
-            [(i, 1) for i in departures] + [(i, 2) for i in arrivals],
-            key=lambda train: datetime.fromisoformat(
-                train[0]["departure"] if train[1] == 1 else train[0]["arrival"]
-            ).timestamp(),
-        )
-        trains = [
-            (i, type)
-            for i, type in trains
-            if datetime.fromisoformat(i["departure"] if type == 1 else i["arrival"]).timestamp() > date.timestamp()  # type: ignore
-        ]
-        brands = self.storage.get_cache("brands") or self.storage.set_cache("brands", self.client.get_brands())
-        parts = []
-        for train, type in trains:
-            time = (
-                f"[bold green]{train['departure'][11:16]}[/bold green]"
-                if type == 1
-                else f"[bold yellow]{train['arrival'][11:16]}[/bold yellow]"
-            )
-            brand = next(iter(i for i in brands if i["id"] == train["brand_id"]), {}).get("logo_text")
-            parts.append(
-                f"{time} [red]{brand}[/red] {train["train_full_name"]}[purple] {train["stations"][0]["name"]} {self.format_position(train["platform"], train["track"])}[/purple]"
-            )
-        self.print("\n".join(parts))
-
-    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 get_train_calendars(self, brand: str, name: str) -> list[TrainCalendar]:
-        brand = brand.upper().strip()
-        name_parts = name.split(" ")
-        if len(name_parts) == 1 and name_parts[0].isnumeric():
-            number = int(name_parts[0])
-            train_name = ""
-        elif len(name) > 1:
-            number = int(name_parts.pop(0))
-            train_name = " ".join(name_parts)
-        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}"
-        try:
-            train_calendars = self.storage.get_cache(cache_id) or self.storage.set_cache(
-                cache_id, self.client.get_train_calendars(brand, number, train_name)
-            )
-        except self.client.errors.KoleoNotFound:
-            self.print(f'[bold red]Train not found: nr={number}, name="{train_name}"[/bold red]')
-            exit(2)
-        return train_calendars["train_calendars"]
-
-    def train_calendar(self, brand: str, name: str):
-        train_calendars = self.get_train_calendars(brand, name)
-        brands = self.storage.get_cache("brands") or self.storage.set_cache("brands", self.client.get_brands())
-        for calendar in train_calendars:
-            brand = next(iter(i for i in brands if i["id"] == calendar["trainBrand"]), {}).get("logo_text", "")
-            parts = [
-                f"[red]{brand}[/red] [bold blue]{calendar['train_nr']}{" "+ v if (v:=calendar.get("train_name")) else ""}[/bold blue]:"
-            ]
-            for k, v in calendar["date_train_map"].items():
-                parts.append(f"  [bold green]{k}[/bold green]: [purple]{v}[/purple]")
-            self.print("\n".join(parts))
-
-    def train_info(self, brand: str, name: str, date: datetime):
-        train_calendars = self.get_train_calendars(brand, name)
-        if not (train_id := train_calendars[0]["date_train_map"].get(date.strftime("%Y-%m-%d"))):
-            self.print(f"[bold red]This train doesn't run on the selected date: {date.strftime("%Y-%m-%d")}[/bold red]")
-            exit(2)
-        self.train_detail(train_id)
-
-    def train_detail(self, train_id: int):
-        train_details = self.client.get_train(train_id)
-        brands = self.storage.get_cache("brands") or self.storage.set_cache("brands", self.client.get_brands())
-        brand = next(iter(i for i in brands if i["id"] == train_details["train"]["brand_id"]), {}).get("logo_text", "")
-
-        parts = [f"[red]{brand}[/red] [bold blue]{train_details["train"]["train_full_name"]}[/bold blue]"]
-        if train_details["train"]["run_desc"]:
-            parts.append(f"  {train_details["train"]["run_desc"]}")
-
-        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"]
-            if stop["vehicle_type"]
-        }
-        if vehicle_types:
-            keys = list(vehicle_types.keys())
-            start = keys[0]
-            for i in range(1, len(keys)):
-                if vehicle_types[keys[i]] != vehicle_types[start]:
-                    parts.append(f"  {start} - {keys[i]}: [bold green]{vehicle_types[start]}[/bold green]")
-                    start = keys[i]
-            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, purchasable: bool, length: int = 1
-    ):
-        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["name"]: i["id"] for i in api_brands}
-        else:
-            connection_brands = {
-                i["name"]: 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.print(f'[bold red]No brands match: "{', '.join(brands)}"[/bold red]')
-                exit(2)
-        results = []
-        fetch_date = date
-        while len(results) < length:
-            connections = self.client.get_connections(
-                start_station["name_slug"],
-                end_station["name_slug"],
-                list(connection_brands.values()),
-                fetch_date,
-                direct,
-                purchasable,
-            )
-            if connections:
-                fetch_date = arr_dep_to_dt(connections[-1]["departure"]) + timedelta(seconds=(30 * 60) + 1)  # wtf
-                results.extend(connections)
-            else:
-                break
-        link = (
-            f"https://koleo.pl/rozklad-pkp/{start_station["name_slug"]}/{end_station["name_slug"]}"
-            + f"/{date.strftime("%d-%m-%Y_%H:%M")}"
-            + f"{"all" if not direct else "direct"}/{"-".join(connection_brands.keys()) if brands else "all"}"
-        )
-        parts = [
-            f"[bold blue][link={link}]{start_station["name"]} → {end_station["name"]} at {date.strftime("%H:%M %d-%m")}[/link][/bold blue]"
-        ]
-        for i in results:
-            arr = arr_dep_to_dt(i["arrival"])
-            dep = arr_dep_to_dt(i["departure"])
-            travel_time = (arr - dep).seconds
-            date_part = f"{arr.strftime("%d-%m")} " if arr.date() != date.date() else ""
-            parts.append(
-                f"[bold green][link=https://koleo.pl/travel-options/{i["id"]}]{date_part}{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 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"])
-                    )
-                )
-                brand = next(iter(i for i in api_brands if i["id"] == train["brand_id"]), {}).get("logo_text")
-                s = f" [red]{brand}[/red] {train["train_full_name"]}[purple] {stop_station['name']} {self.format_position(stop["platform"], stop["track"])}[/purple]"
-                parts[-1] += (
-                    f" [red]{brand}[/red] {train["train_full_name"]}[purple] {stop_station['name']} {self.format_position(stop["platform"], stop["track"])}[/purple]"
-                )
-                for constriction in i["constriction_info"]:
-                    parts.append(f" [bold red]- {constriction}[/bold red]")
-            else:
-                for constriction in i["constriction_info"]:
-                    parts.append(f" [bold red]- {constriction}[/bold red]")
-                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_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']} {self.format_position(fs["platform"], fs["track"])}[/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_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']} {self.format_position(ls["platform"], ls["track"])}[/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 = []
-        brands = self.storage.get_cache("brands") or self.storage.set_cache("brands", self.client.get_brands())
-        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("logo_text")
-            parts.append(
-                f"[bold green]{time[11:16]}[/bold green] [red]{brand}[/red] {train["train_full_name"]}[purple] {train["stations"][0]["name"]} {self.format_position(train["platform"], train["track"])}[/purple]"
-            )
-        return "\n".join(parts)
-
-    def train_route_table(self, train: TrainDetailResponse):
-        parts = []
-        for stop in train["stops"]:
-            arr = arr_dep_to_dt(stop["arrival"])
-            dep = arr_dep_to_dt(stop["departure"])
-            parts.append(
-                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"]} {self.format_position(stop["platform"])} [/purple]"
-            )
-        return "\n".join(parts)
-
-    def format_position(self, platform: str, track: str | None = None):
-        res = str(convert_platform_number(platform) or "" if not self.storage.use_roman_numerals else platform)
-        if track is not None and track != "":
-            res += f"/{track}"
-        return res
-
-    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.print(f'[bold red]Station not found: "{station}"[/bold red]')
-            exit(2)
-
-
-def main():
-    cli = CLI()
-
-    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 and formatting", action="store_true", default=False)
-    subparsers = parser.add_subparsers(title="actions", required=False)  # type: ignore
-
-    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",
-        default=None,
-        nargs="*",
-        action=RemainderString,
-    )
-    departures.add_argument(
-        "-d",
-        "--date",
-        help="the departure date",
-        type=lambda s: parse_datetime(s),
-        default=datetime.now(),
-    )
-    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", "przyjazdy", "p"], help="Allows you to list station departures"
-    )
-    arrivals.add_argument(
-        "station",
-        help="The station name",
-        default=None,
-        nargs="*",
-        action=RemainderString,
-    )
-    arrivals.add_argument(
-        "-d",
-        "--date",
-        help="the arrival date",
-        type=lambda s: parse_datetime(s),
-        default=datetime.now(),
-    )
-    arrivals.add_argument("-s", "--save", help="save the station as your default one", action="store_true")
-    arrivals.set_defaults(func=cli.full_arrivals, pass_=["station", "date"])
-
-    all_trains = subparsers.add_parser(
-        "all", aliases=["w", "wszystkie", "all_trains", "pociagi"], help="Allows you to list all station trains"
-    )
-    all_trains.add_argument(
-        "station",
-        help="The station name",
-        default=None,
-        nargs="*",
-        action=RemainderString,
-    )
-    all_trains.add_argument(
-        "-d",
-        "--date",
-        help="the date",
-        type=lambda s: parse_datetime(s),
-        default=datetime.now(),
-    )
-    all_trains.add_argument("-s", "--save", help="save the station as your default one", action="store_true")
-    all_trains.set_defaults(func=cli.all_trains, pass_=["station", "date"])
-
-    train_route = subparsers.add_parser(
-        "trainroute",
-        aliases=["r", "tr", "t", "poc", "pociąg"],
-        help="Allows you to check the train's route",
-    )
-    train_route.add_argument("brand", help="The brand name", type=str)
-    train_route.add_argument("name", help="The train name", nargs="+", action=RemainderString)
-    train_route.add_argument(
-        "-d",
-        "--date",
-        help="the date",
-        type=lambda s: parse_datetime(s),
-        default=datetime.now(),
-    )
-    train_route.set_defaults(func=cli.train_info, pass_=["brand", "name", "date"])
-
-    train_calendar = subparsers.add_parser(
-        "traincalendar",
-        aliases=["kursowanie", "tc", "k"],
-        help="Allows you to check what days the train runs on",
-    )
-    train_calendar.add_argument("brand", help="The brand name", type=str)
-    train_calendar.add_argument("name", help="The train name", nargs="+", action=RemainderString)
-    train_calendar.set_defaults(func=cli.train_calendar, pass_=["brand", "name"])
-
-    train_detail = subparsers.add_parser(
-        "traindetail",
-        aliases=["td", "tid", "id", "idpoc"],
-        help="Allows you to show the train's route given it's koleo ID",
-    )
-    train_detail.add_argument("train_id", help="The koleo ID", type=int)
-    train_detail.set_defaults(func=cli.train_detail, pass_=["train_id"])
-
-    stations = subparsers.add_parser(
-        "stations", aliases=["s", "find", "f", "stacje", "ls", "q"], help="Allows you to find stations by their name"
-    )
-    stations.add_argument(
-        "query",
-        help="The station name",
-        default=None,
-        nargs="*",
-        action=RemainderString,
-    )
-    stations.set_defaults(func=cli.find_station, pass_=["query"])
-
-    connections = subparsers.add_parser(
-        "connections",
-        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)
-    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.add_argument(
-        "-l",
-        "--length",
-        help="fetch at least n connections",
-        type=int,
-        default=1,
-    )
-    connections.set_defaults(
-        func=cli.connections, pass_=["start", "end", "brands", "date", "direct", "purchasable", "length"]
-    )
-
-    args = parser.parse_args()
-
-    storage = Storage.load(path=args.config)
-    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:
-        storage.favourite_station = args.station
-    if not hasattr(args, "func"):
-        if storage.favourite_station:
-            cli.full_departures(storage.favourite_station, datetime.now())
-        else:
-            parser.print_help()
-            exit()
-    else:
-        args.func(**{k: v for k, v in args.__dict__.items() if k in getattr(args, "pass_", [])})
-
-    if storage._dirty:
-        storage.save()