Split travel YAML parsing into function.
This commit is contained in:
parent
26aa8a8ad8
commit
0181b9ce66
|
@ -1,6 +1,7 @@
|
|||
"""Travel."""
|
||||
|
||||
import os
|
||||
import typing
|
||||
|
||||
import yaml
|
||||
|
||||
|
@ -8,10 +9,17 @@ from .types import Event
|
|||
|
||||
Leg = dict[str, str]
|
||||
|
||||
TravelList = list[dict[str, typing.Any]]
|
||||
|
||||
|
||||
def parse_yaml(travel_type: str, data_dir: str) -> TravelList:
|
||||
"""Parse flights YAML and return list of travel."""
|
||||
filepath = os.path.join(data_dir, travel_type + ".yaml")
|
||||
return typing.cast(TravelList, yaml.safe_load(open(filepath)))
|
||||
|
||||
|
||||
def get_flights(data_dir: str) -> list[Event]:
|
||||
"""Get travel events."""
|
||||
filepath = os.path.join(data_dir, "flights.yaml")
|
||||
return [
|
||||
Event(
|
||||
date=item["depart"],
|
||||
|
@ -20,7 +28,7 @@ def get_flights(data_dir: str) -> list[Event]:
|
|||
title=f'flight from {item["from"]} to {item["to"]} ({flight_number(item)})',
|
||||
url=item.get("url"),
|
||||
)
|
||||
for item in yaml.safe_load(open(filepath))
|
||||
for item in parse_yaml("flights", data_dir)
|
||||
if item["depart"].date()
|
||||
]
|
||||
|
||||
|
@ -28,8 +36,7 @@ def get_flights(data_dir: str) -> list[Event]:
|
|||
def get_trains(data_dir: str) -> list[Event]:
|
||||
"""Get train events."""
|
||||
events: list[Event] = []
|
||||
filepath = os.path.join(data_dir, "trains.yaml")
|
||||
for item in yaml.safe_load(open(filepath)):
|
||||
for item in parse_yaml("trains", data_dir):
|
||||
events += [
|
||||
Event(
|
||||
date=leg["depart"],
|
||||
|
|
Loading…
Reference in a new issue