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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of koleo-cli might be problematic. Click here for more details.

koleo/cli.py DELETED
@@ -1,608 +0,0 @@
1
- import re
2
- from argparse import ArgumentParser
3
- from datetime import datetime, timedelta
4
-
5
- from rich.console import Console
6
-
7
- from .api import KoleoAPI
8
- from .storage import DEFAULT_CONFIG_PATH, Storage
9
- from .types import ExtendedBaseStationInfo, TrainCalendar, TrainOnStationInfo, TrainStop
10
- from .utils import RemainderString, arr_dep_to_dt, convert_platform_number, name_to_slug, parse_datetime
11
-
12
-
13
- class CLI:
14
- def __init__(
15
- self,
16
- no_color: bool = False,
17
- client: KoleoAPI | None = None,
18
- storage: Storage | None = None,
19
- ) -> None:
20
- self._client = client
21
- self._storage = storage
22
- self.no_color = no_color
23
- self.console = Console(color_system="standard", no_color=no_color, highlight=False)
24
-
25
- def print(self, text, *args, **kwargs):
26
- if self.no_color:
27
- result = re.sub(r"\[[^\]]*\]", "", text)
28
- print(result)
29
- else:
30
- self.console.print(text, *args, **kwargs)
31
-
32
- @property
33
- def client(self) -> KoleoAPI:
34
- if not self._client:
35
- raise ValueError("Client not set!")
36
- return self._client
37
-
38
- @client.setter
39
- def client(self, client: KoleoAPI):
40
- self._client = client
41
-
42
- @property
43
- def storage(self) -> Storage:
44
- if not self._storage:
45
- raise ValueError("Storage not set!")
46
- return self._storage
47
-
48
- @storage.setter
49
- def storage(self, storage: Storage):
50
- self._storage = storage
51
-
52
- def get_departures(self, station_id: int, date: datetime):
53
- cache_id = f"dep-{station_id}-{date.strftime("%Y-%m-%d")}"
54
- trains = self.storage.get_cache(cache_id) or self.storage.set_cache(
55
- cache_id, self.client.get_departures(station_id, date)
56
- )
57
- trains = [
58
- i
59
- for i in trains
60
- if datetime.fromisoformat(i["departure"]).timestamp() > date.timestamp() # type: ignore
61
- ]
62
- table = self.trains_on_station_table(trains)
63
- self.print(table)
64
- return table
65
-
66
- def get_arrivals(self, station_id: int, date: datetime):
67
- cache_id = f"arr-{station_id}-{date.strftime("%Y-%m-%d")}"
68
- trains = self.storage.get_cache(cache_id) or self.storage.set_cache(
69
- cache_id, self.client.get_arrivals(station_id, date)
70
- )
71
- trains = [
72
- i
73
- for i in trains
74
- if datetime.fromisoformat(i["arrival"]).timestamp() > date.timestamp() # type: ignore
75
- ]
76
- table = self.trains_on_station_table(trains, type=2)
77
- self.print(table)
78
- return table
79
-
80
- def full_departures(self, station: str, date: datetime):
81
- st = self.get_station(station)
82
- station_info = f"[bold blue][link=https://koleo.pl/dworzec-pkp/{st["name_slug"]}/odjazdy/{date.strftime("%Y-%m-%d")}]{st["name"]} at {date.strftime("%d-%m %H:%M")}[/bold blue] ID: {st["id"]}[/link]"
83
- self.print(station_info)
84
- self.get_departures(st["id"], date)
85
-
86
- def full_arrivals(self, station: str, date: datetime):
87
- st = self.get_station(station)
88
- station_info = f"[bold blue][link=https://koleo.pl/dworzec-pkp/{st["name_slug"]}/przyjazdy/{date.strftime("%Y-%m-%d")}]{st["name"]} at {date.strftime("%d-%m %H:%M")}[/bold blue] ID: {st["id"]}[/link]"
89
- self.print(station_info)
90
- self.get_arrivals(st["id"], date)
91
-
92
- def all_trains(self, station: str, date: datetime):
93
- st = self.get_station(station)
94
- station_info = f"[bold blue][link=https://koleo.pl/dworzec-pkp/{st["name_slug"]}/odjazdy/{date.strftime("%Y-%m-%d")}]{st["name"]} at {date.strftime("%d-%m %H:%M")}[/bold blue] ID: {st["id"]}[/link]"
95
- self.print(station_info)
96
- arr_cache_id = f"arr-{st['id']}-{date.strftime("%Y-%m-%d")}"
97
- dep_cache_id = f"dep-{st['id']}-{date.strftime("%Y-%m-%d")}"
98
- arrivals = self.storage.get_cache(arr_cache_id) or self.storage.set_cache(
99
- arr_cache_id, self.client.get_arrivals(st["id"], date)
100
- )
101
- departures = self.storage.get_cache(dep_cache_id) or self.storage.set_cache(
102
- dep_cache_id, self.client.get_departures(st["id"], date)
103
- )
104
- trains = sorted(
105
- [(i, 1) for i in departures] + [(i, 2) for i in arrivals],
106
- key=lambda train: (
107
- datetime.fromisoformat(train[0]["departure"]) + timedelta(microseconds=1)
108
- if train[1] == 1
109
- else (datetime.fromisoformat(train[0]["arrival"]))
110
- ).timestamp(),
111
- )
112
- trains = [
113
- (i, type)
114
- for i, type in trains
115
- if datetime.fromisoformat(i["departure"] if type == 1 else i["arrival"]).timestamp() > date.timestamp() # type: ignore
116
- ]
117
- brands = self.storage.get_cache("brands") or self.storage.set_cache("brands", self.client.get_brands())
118
- parts = []
119
- for train, type in trains:
120
- time = (
121
- f"[bold green]{train['departure'][11:16]}[/bold green]"
122
- if type == 1
123
- else f"[bold yellow]{train['arrival'][11:16]}[/bold yellow]"
124
- )
125
- brand = next(iter(i for i in brands if i["id"] == train["brand_id"]), {}).get("logo_text")
126
- parts.append(
127
- f"{time} [red]{brand}[/red] {train["train_full_name"]}[purple] {train["stations"][0]["name"]} {self.format_position(train["platform"], train["track"])}[/purple]"
128
- )
129
- self.print("\n".join(parts))
130
-
131
- def find_station(self, query: str | None):
132
- if query:
133
- stations = self.client.find_station(query)
134
- else:
135
- stations = self.storage.get_cache("stations") or self.storage.set_cache(
136
- "stations", self.client.get_stations()
137
- )
138
- for st in stations:
139
- self.print(
140
- f"[bold blue][link=https://koleo.pl/dworzec-pkp/{st["name_slug"]}]{st["name"]}[/bold blue] ID: {st["id"]}[/link]"
141
- )
142
-
143
- def get_train_calendars(self, brand: str, name: str) -> list[TrainCalendar]:
144
- brand = brand.upper().strip()
145
- name_parts = name.split(" ")
146
- if len(name_parts) == 1 and name_parts[0].isnumeric():
147
- number = int(name_parts[0])
148
- train_name = ""
149
- elif len(name) > 1:
150
- number = int(name_parts.pop(0))
151
- train_name = " ".join(name_parts)
152
- else:
153
- raise ValueError("Invalid train name!")
154
- brands = self.storage.get_cache("brands") or self.storage.set_cache("brands", self.client.get_brands())
155
- if brand not in [i["name"] for i in brands]:
156
- res = {i["logo_text"]: i["name"] for i in brands}.get(brand)
157
- if not res:
158
- raise ValueError("Invalid brand name!")
159
- brand = res
160
- cache_id = f"tc-{brand}-{number}-{name}"
161
- try:
162
- train_calendars = self.storage.get_cache(cache_id) or self.storage.set_cache(
163
- cache_id, self.client.get_train_calendars(brand, number, train_name)
164
- )
165
- except self.client.errors.KoleoNotFound:
166
- self.print(f'[bold red]Train not found: nr={number}, name="{train_name}"[/bold red]')
167
- exit(2)
168
- return train_calendars["train_calendars"]
169
-
170
- def train_calendar(self, brand: str, name: str):
171
- train_calendars = self.get_train_calendars(brand, name)
172
- brands = self.storage.get_cache("brands") or self.storage.set_cache("brands", self.client.get_brands())
173
- for calendar in train_calendars:
174
- brand_obj = next(iter(i for i in brands if i["id"] == calendar["trainBrand"]), {})
175
- link = f"https://koleo.pl/pociag/{brand_obj["name"]}/{name.replace(" ", "-")}"
176
- brand = brand_obj.get("logo_text", "")
177
- parts = [
178
- f"[red][link={link}]{brand}[/red] [bold blue]{calendar['train_nr']}{" "+ v if (v:=calendar.get("train_name")) else ""}[/bold blue]:[/link]"
179
- ]
180
- for k, v in sorted(calendar["date_train_map"].items(), key=lambda x: datetime.strptime(x[0], "%Y-%m-%d")):
181
- parts.append(f" [bold green]{k}[/bold green]: [purple]{v}[/purple]")
182
- self.print("\n".join(parts))
183
-
184
- def train_info(
185
- self, brand: str, name: str, date: datetime, closest: bool, show_stations: tuple[str, str] | None = None
186
- ):
187
- train_calendars = self.get_train_calendars(brand, name)
188
- if closest:
189
- dates = sorted([datetime.strptime(i, "%Y-%m-%d") for i in train_calendars[0]["dates"]])
190
- date = next(iter(i for i in dates if i > date)) or next(iter(i for i in reversed(dates) if i < date))
191
- if not (train_id := train_calendars[0]["date_train_map"].get(date.strftime("%Y-%m-%d"))):
192
- self.print(f"[bold red]This train doesn't run on the selected date: {date.strftime("%Y-%m-%d")}[/bold red]")
193
- exit(2)
194
- self.train_detail(train_id, date=date.strftime("%Y-%m-%d"), show_stations=show_stations)
195
-
196
- def train_detail(self, train_id: int, date: str | None = None, show_stations: tuple[str, str] | None = None):
197
- train_details = self.client.get_train(train_id)
198
- brands = self.storage.get_cache("brands") or self.storage.set_cache("brands", self.client.get_brands())
199
- brand_obj = next(iter(i for i in brands if i["id"] == train_details["train"]["brand_id"]), {})
200
- brand = brand_obj.get("logo_text", "")
201
- parts = [f"[red]{brand}[/red] [bold blue]{train_details["train"]["train_full_name"]}[/bold blue]"]
202
-
203
- link = (
204
- f"https://koleo.pl/pociag/{brand_obj["name"]}/{train_details["train"]["train_full_name"].replace(" ", "-")}"
205
- )
206
- if date:
207
- link += f"/{date}"
208
- parts[0] = f"[link={link}]{parts[0]}[/link]"
209
-
210
- if train_details["train"]["run_desc"]:
211
- parts.append(f" {train_details["train"]["run_desc"]}")
212
-
213
- if show_stations:
214
- slugs = [self.get_station(i)["name_slug"] for i in show_stations]
215
- first_stop, fs_index = next(
216
- iter((i, n) for n, i in enumerate(train_details["stops"]) if i["station_slug"] == slugs[0]), {}
217
- )
218
- last_stop, ls_index = next(
219
- iter((i, n + 1) for n, i in enumerate(train_details["stops"]) if i["station_slug"] == slugs[1]), {}
220
- )
221
- if fs_index >= ls_index:
222
- self.print("[bold red]Station B has to be after station A (-s / --show_stations)[/bold red]")
223
- exit(2)
224
- else:
225
- first_stop, last_stop = train_details["stops"][0], train_details["stops"][-1]
226
- fs_index, ls_index = 0, len(train_details["stops"]) + 1
227
- route_start = arr_dep_to_dt(first_stop["departure"])
228
- route_end = arr_dep_to_dt(last_stop["arrival"])
229
-
230
- if route_end.hour < route_start.hour or (
231
- route_end.hour == route_start.hour and route_end.minute < route_end.minute
232
- ):
233
- route_end += timedelta(days=1)
234
-
235
- travel_time = route_end - route_start
236
- speed = (last_stop["distance"] - first_stop["distance"]) / 1000 / travel_time.seconds * 3600
237
- parts.append(
238
- f"[white] {travel_time.seconds//3600}h{(travel_time.seconds % 3600)/60:.0f}m {speed:^4.1f}km/h [/white]"
239
- )
240
-
241
- vehicle_types: dict[str, str] = {
242
- stop["station_display_name"]: stop["vehicle_type"]
243
- for stop in train_details["stops"]
244
- if stop["vehicle_type"]
245
- }
246
- if vehicle_types:
247
- keys = list(vehicle_types.keys())
248
- start = keys[0]
249
- for i in range(1, len(keys)):
250
- if vehicle_types[keys[i]] != vehicle_types[start]:
251
- parts.append(f" {start} - {keys[i]}: [bold green]{vehicle_types[start]}[/bold green]")
252
- start = keys[i]
253
- parts.append(f" {start} - {keys[-1]}: [bold green]{vehicle_types[start]}[/bold green]")
254
- self.print("\n".join(parts))
255
- self.print(self.train_route_table(train_details["stops"][fs_index:ls_index]))
256
-
257
- def connections(
258
- self, start: str, end: str, date: datetime, brands: list[str], direct: bool, purchasable: bool, length: int = 1
259
- ):
260
- start_station = self.get_station(start)
261
- end_station = self.get_station(end)
262
- brands = [i.lower().strip() for i in brands]
263
- api_brands = self.storage.get_cache("brands") or self.storage.set_cache("brands", self.client.get_brands())
264
- if not brands:
265
- connection_brands = {i["name"]: i["id"] for i in api_brands}
266
- else:
267
- connection_brands = {
268
- i["name"]: i["id"]
269
- for i in api_brands
270
- if i["name"].lower().strip() in brands or i["logo_text"].lower().strip() in brands
271
- }
272
- if not connection_brands:
273
- self.print(f'[bold red]No brands match: "{', '.join(brands)}"[/bold red]')
274
- exit(2)
275
- results = []
276
- fetch_date = date
277
- while len(results) < length:
278
- connections = self.client.get_connections(
279
- start_station["name_slug"],
280
- end_station["name_slug"],
281
- list(connection_brands.values()),
282
- fetch_date,
283
- direct,
284
- purchasable,
285
- )
286
- if connections:
287
- fetch_date = arr_dep_to_dt(connections[-1]["departure"]) + timedelta(seconds=(30 * 60) + 1) # wtf
288
- results.extend(connections)
289
- else:
290
- break
291
- link = (
292
- f"https://koleo.pl/rozklad-pkp/{start_station["name_slug"]}/{end_station["name_slug"]}"
293
- + f"/{date.strftime("%d-%m-%Y_%H:%M")}"
294
- + f"/{"all" if not direct else "direct"}/{"-".join(connection_brands.keys()) if brands else "all"}"
295
- )
296
- parts = [
297
- f"[bold blue][link={link}]{start_station["name"]} → {end_station["name"]} at {date.strftime("%H:%M %d-%m")}[/link][/bold blue]"
298
- ]
299
-
300
- for i in results:
301
- arr = arr_dep_to_dt(i["arrival"])
302
- dep = arr_dep_to_dt(i["departure"])
303
- travel_time = (arr - dep).seconds
304
- date_part = f"{arr.strftime("%d-%m")} " if arr.date() != date.date() else ""
305
- parts.append(
306
- 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]"
307
- )
308
- if len(i["trains"]) == 1:
309
- train = i["trains"][0]
310
- brand = next(iter(i for i in api_brands if i["id"] == train["brand_id"]), {}).get("logo_text")
311
-
312
- fs = next(iter(i for i in train["stops"] if i["station_id"] == train["start_station_id"]), {})
313
- fs_station = (
314
- start_station
315
- if fs["station_id"] == start_station["id"]
316
- else self.storage.get_cache(f"st-{fs['station_id']}")
317
- or self.storage.set_cache(f"st-{fs['station_id']}", self.client.get_station_by_id(fs["station_id"]))
318
- )
319
-
320
- ls = next(iter(i for i in train["stops"] if i["station_id"] == train["end_station_id"]), {})
321
- ls_station = (
322
- start_station
323
- if ls["station_id"] == start_station["id"]
324
- else self.storage.get_cache(f"st-{ls['station_id']}")
325
- or self.storage.set_cache(f"st-{ls['station_id']}", self.client.get_station_by_id(ls["station_id"]))
326
- )
327
-
328
- parts[-1] += (
329
- f" [red]{brand}[/red] {train["train_full_name"]}[purple] {fs_station['name']} {self.format_position(fs["platform"], fs["track"])}[/purple] - [purple]{ls_station['name']} {self.format_position(ls["platform"], ls["track"])}[/purple]"
330
- )
331
- for constriction in i["constriction_info"]:
332
- parts.append(f" [bold red]- {constriction}[/bold red]")
333
- else:
334
- for constriction in i["constriction_info"]:
335
- parts.append(f" [bold red]- {constriction}[/bold red]")
336
- previous_arrival: datetime | None = None
337
- for train in i["trains"]:
338
- brand = next(iter(i for i in api_brands if i["id"] == train["brand_id"]), {}).get("logo_text")
339
-
340
- # first stop
341
-
342
- fs = next(iter(i for i in train["stops"] if i["station_id"] == train["start_station_id"]), {})
343
- fs_station = (
344
- start_station
345
- if fs["station_id"] == start_station["id"]
346
- else self.storage.get_cache(f"st-{fs['station_id']}")
347
- or self.storage.set_cache(
348
- f"st-{fs['station_id']}", self.client.get_station_by_id(fs["station_id"])
349
- )
350
- )
351
- # fs_arr = arr_dep_to_dt(fs["arrival"])
352
- fs_dep = arr_dep_to_dt(fs["departure"])
353
- fs_info = f"[bold green]{fs_dep.strftime("%H:%M")} [/bold green][purple]{fs_station['name']} {self.format_position(fs["platform"], fs["track"])}[/purple]"
354
-
355
- # last stop
356
-
357
- ls = next(iter(i for i in train["stops"] if i["station_id"] == train["end_station_id"]), {})
358
- ls_station = (
359
- start_station
360
- if ls["station_id"] == start_station["id"]
361
- else self.storage.get_cache(f"st-{ls['station_id']}")
362
- or self.storage.set_cache(
363
- f"st-{ls['station_id']}", self.client.get_station_by_id(ls["station_id"])
364
- )
365
- )
366
- ls_arr = arr_dep_to_dt(ls["arrival"])
367
- # ls_dep = arr_dep_to_dt(ls["departure"])
368
- ls_info = f"[bold green]{ls_arr.strftime("%H:%M")} [/bold green][purple]{ls_station['name']} {self.format_position(ls["platform"], ls["track"])}[/purple]"
369
- connection_time = (fs_dep - previous_arrival).seconds if previous_arrival else ""
370
- previous_arrival = ls_arr
371
- if connection_time:
372
- parts.append(
373
- f" {connection_time//3600}h{(connection_time % 3600)/60:.0f}m at [purple]{fs_station['name']}[/purple]"
374
- )
375
- parts.append(f" [red]{brand}[/red] {train["train_full_name"]} {fs_info} - {ls_info}")
376
- self.print("\n".join(parts))
377
-
378
- def trains_on_station_table(self, trains: list[TrainOnStationInfo], type: int = 1):
379
- parts = []
380
- brands = self.storage.get_cache("brands") or self.storage.set_cache("brands", self.client.get_brands())
381
- for train in trains:
382
- time, color = (train["departure"], "green") if type == 1 else (train["arrival"], "yellow")
383
- assert time
384
- brand = next(iter(i for i in brands if i["id"] == train["brand_id"]), {}).get("logo_text")
385
- parts.append(
386
- f"[bold {color}]{time[11:16]}[/bold {color}] [red]{brand}[/red] {train["train_full_name"]}[purple] {train["stations"][0]["name"]} {self.format_position(train["platform"], train["track"])}[/purple]"
387
- )
388
- return "\n".join(parts)
389
-
390
- def train_route_table(self, stops: list[TrainStop]):
391
- parts = []
392
- last_real_distance = stops[0]["distance"]
393
- for stop in stops:
394
- arr = arr_dep_to_dt(stop["arrival"])
395
- dep = arr_dep_to_dt(stop["departure"])
396
- distance = stop["distance"] - last_real_distance
397
- parts.append(
398
- f"[white underline]{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]"
399
- )
400
- return "\n".join(parts)
401
-
402
- def format_position(self, platform: str, track: str | None = None):
403
- res = str(convert_platform_number(platform) or "" if not self.storage.use_roman_numerals else platform)
404
- if track is not None and track != "":
405
- res += f"/{track}"
406
- return res
407
-
408
- def get_station(self, station: str) -> ExtendedBaseStationInfo:
409
- slug = name_to_slug(station)
410
- try:
411
- return self.storage.get_cache(f"st-{slug}") or self.storage.set_cache(
412
- f"st-{slug}", self.client.get_station_by_slug(slug)
413
- )
414
- except self.client.errors.KoleoNotFound:
415
- self.print(f'[bold red]Station not found: "{station}"[/bold red]')
416
- exit(2)
417
-
418
-
419
- def main():
420
- cli = CLI()
421
-
422
- parser = ArgumentParser("koleo", description="Koleo CLI")
423
- parser.add_argument("-c", "--config", help="Custom config path.", default=DEFAULT_CONFIG_PATH)
424
- parser.add_argument("--nocolor", help="Disable color output and formatting", action="store_true", default=False)
425
- subparsers = parser.add_subparsers(title="actions", required=False) # type: ignore
426
-
427
- departures = subparsers.add_parser(
428
- "departures", aliases=["d", "dep", "odjazdy", "o"], help="Allows you to list station departures"
429
- )
430
- departures.add_argument(
431
- "station",
432
- help="The station name",
433
- default=None,
434
- nargs="*",
435
- action=RemainderString,
436
- )
437
- departures.add_argument(
438
- "-d",
439
- "--date",
440
- help="the departure date",
441
- type=lambda s: parse_datetime(s),
442
- default=datetime.now(),
443
- )
444
- departures.add_argument("-s", "--save", help="save the station as your default one", action="store_true")
445
- departures.set_defaults(func=cli.full_departures, pass_=["station", "date"])
446
-
447
- arrivals = subparsers.add_parser(
448
- "arrivals", aliases=["a", "arr", "przyjazdy", "p"], help="Allows you to list station departures"
449
- )
450
- arrivals.add_argument(
451
- "station",
452
- help="The station name",
453
- default=None,
454
- nargs="*",
455
- action=RemainderString,
456
- )
457
- arrivals.add_argument(
458
- "-d",
459
- "--date",
460
- help="the arrival date",
461
- type=lambda s: parse_datetime(s),
462
- default=datetime.now(),
463
- )
464
- arrivals.add_argument("-s", "--save", help="save the station as your default one", action="store_true")
465
- arrivals.set_defaults(func=cli.full_arrivals, pass_=["station", "date"])
466
-
467
- all_trains = subparsers.add_parser(
468
- "all", aliases=["w", "wszystkie", "all_trains", "pociagi"], help="Allows you to list all station trains"
469
- )
470
- all_trains.add_argument(
471
- "station",
472
- help="The station name",
473
- default=None,
474
- nargs="*",
475
- action=RemainderString,
476
- )
477
- all_trains.add_argument(
478
- "-d",
479
- "--date",
480
- help="the date",
481
- type=lambda s: parse_datetime(s),
482
- default=datetime.now(),
483
- )
484
- all_trains.add_argument("-s", "--save", help="save the station as your default one", action="store_true")
485
- all_trains.set_defaults(func=cli.all_trains, pass_=["station", "date"])
486
-
487
- train_route = subparsers.add_parser(
488
- "trainroute",
489
- aliases=["r", "tr", "t", "poc", "pociąg"],
490
- help="Allows you to check the train's route",
491
- )
492
- train_route.add_argument("brand", help="The brand name", type=str)
493
- train_route.add_argument("name", help="The train name", nargs="+", action=RemainderString)
494
- train_route.add_argument(
495
- "-d",
496
- "--date",
497
- help="the date",
498
- type=lambda s: parse_datetime(s),
499
- default=datetime.now(),
500
- )
501
- train_route.add_argument(
502
- "-c",
503
- "--closest",
504
- help="ignores date, fetches closest date from the train calendar",
505
- action="store_true",
506
- default=False,
507
- )
508
- train_route.add_argument(
509
- "-s", "--show_stations", help="limit the result to A->B", action="extend", nargs=2, type=str, default=None
510
- )
511
- train_route.set_defaults(func=cli.train_info, pass_=["brand", "name", "date", "closest", "show_stations"])
512
-
513
- train_calendar = subparsers.add_parser(
514
- "traincalendar",
515
- aliases=["kursowanie", "tc", "k"],
516
- help="Allows you to check what days the train runs on",
517
- )
518
- train_calendar.add_argument("brand", help="The brand name", type=str)
519
- train_calendar.add_argument("name", help="The train name", nargs="+", action=RemainderString)
520
- train_calendar.set_defaults(func=cli.train_calendar, pass_=["brand", "name"])
521
-
522
- train_detail = subparsers.add_parser(
523
- "traindetail",
524
- aliases=["td", "tid", "id", "idpoc"],
525
- help="Allows you to show the train's route given it's koleo ID",
526
- )
527
- train_detail.add_argument(
528
- "-s", "--show_stations", help="limit the result to A->B", action="extend", nargs=2, type=str, default=None
529
- )
530
- train_detail.add_argument("train_id", help="The koleo ID", type=int)
531
- train_detail.set_defaults(func=cli.train_detail, pass_=["train_id", "show_stations"])
532
-
533
- stations = subparsers.add_parser(
534
- "stations", aliases=["s", "find", "f", "stacje", "ls", "q"], help="Allows you to find stations by their name"
535
- )
536
- stations.add_argument(
537
- "query",
538
- help="The station name",
539
- default=None,
540
- nargs="*",
541
- action=RemainderString,
542
- )
543
- stations.set_defaults(func=cli.find_station, pass_=["query"])
544
-
545
- connections = subparsers.add_parser(
546
- "connections",
547
- aliases=["do", "z", "szukaj", "path"],
548
- help="Allows you to search for connections from a to b",
549
- )
550
- connections.add_argument("start", help="The starting station", type=str)
551
- connections.add_argument("end", help="The end station", type=str)
552
- connections.add_argument(
553
- "-d",
554
- "--date",
555
- help="the date",
556
- type=lambda s: parse_datetime(s),
557
- default=datetime.now(),
558
- )
559
- connections.add_argument(
560
- "-b", "--brands", help="Brands to include", action="extend", nargs="+", type=str, default=[]
561
- )
562
- connections.add_argument(
563
- "-f",
564
- "--direct",
565
- help="whether or not the result should only include direct trains",
566
- action="store_true",
567
- default=False,
568
- )
569
- connections.add_argument(
570
- "-p",
571
- "--purchasable",
572
- help="whether or not the result should only trains purchasable on koleo",
573
- action="store_true",
574
- default=False,
575
- )
576
- connections.add_argument(
577
- "-l",
578
- "--length",
579
- help="fetch at least n connections",
580
- type=int,
581
- default=1,
582
- )
583
- connections.set_defaults(
584
- func=cli.connections, pass_=["start", "end", "brands", "date", "direct", "purchasable", "length"]
585
- )
586
-
587
- args = parser.parse_args()
588
-
589
- storage = Storage.load(path=args.config)
590
- client = KoleoAPI()
591
- cli.client, cli.storage = client, storage
592
- cli.console.no_color = args.nocolor
593
- cli.no_color = args.nocolor
594
- if hasattr(args, "station") and args.station is None:
595
- args.station = storage.favourite_station
596
- elif hasattr(args, "station") and args.save:
597
- storage.favourite_station = args.station
598
- if not hasattr(args, "func"):
599
- if storage.favourite_station:
600
- cli.full_departures(storage.favourite_station, datetime.now())
601
- else:
602
- parser.print_help()
603
- exit()
604
- else:
605
- args.func(**{k: v for k, v in args.__dict__.items() if k in getattr(args, "pass_", [])})
606
-
607
- if storage.dirty:
608
- storage.save()