From 0181b9ce66cd9d67eb40c2433c241be4af7cea25 Mon Sep 17 00:00:00 2001
From: Edward Betts <edward@4angle.com>
Date: Sun, 19 Nov 2023 08:45:34 -0300
Subject: [PATCH] Split travel YAML parsing into function.
---
agenda/travel.py | 15 +++++++++++----
1 file changed, 11 insertions(+), 4 deletions(-)
diff --git a/agenda/travel.py b/agenda/travel.py
index 302e2c0..bc1771d 100644
--- a/agenda/travel.py
+++ b/agenda/travel.py
@@ -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"],