Split travel code into own file

This commit is contained in:
Edward Betts 2023-11-05 12:52:44 +00:00
parent f7022c992d
commit ef45d9c508
2 changed files with 69 additions and 52 deletions

View file

@ -23,7 +23,17 @@ from dateutil.relativedelta import FR, relativedelta
from agenda import thespacedevs
from . import birthday, calendar, economist, fx, gwr, markets, sun, waste_schedule
from . import (
birthday,
calendar,
economist,
fx,
gwr,
markets,
sun,
travel,
waste_schedule,
)
from .types import Event
warnings.simplefilter(action="ignore", category=FutureWarning)
@ -299,56 +309,6 @@ def get_accommodation(from_date: date, filepath: str) -> list[Event]:
return events
Leg = dict[str, str]
def get_travel(
from_date: date,
method: str,
filepath: str,
extra: typing.Callable[[Leg], str] | None = None,
) -> list[Event]:
"""Get travel events."""
def title(item: Leg) -> str:
ret = f'{method} from {item["from"]} to {item["to"]}'
if extra:
ret += f" ({extra(item)})"
return ret
return [
Event(
date=item["depart"],
end_date=item["arrive"],
name="transport",
title=title(item),
url=item.get("url"),
)
for item in yaml.safe_load(open(filepath))
if item["depart"].date() >= from_date
]
def flight_number(flight: Leg) -> str:
"""Flight number."""
airline_code = flight["airline"]
# make sure this is the airline code, not the airline name
assert " " not in airline_code and not any(c.islower() for c in airline_code)
return airline_code + flight["flight_number"]
def get_all_travel_events(from_date: date) -> list[Event]:
"""Get all flights and rail journeys."""
data_dir = config["data"]["personal-data"]
trains = get_travel(from_date, "train", os.path.join(data_dir, "trains.yaml"))
flights = get_travel(
from_date, "flight", os.path.join(data_dir, "flights.yaml"), extra=flight_number
)
return trains + flights
def waste_collection_events() -> list[Event]:
"""Waste colllection events."""
postcode = "BS48 3HG"
@ -420,7 +380,7 @@ def get_data(now: datetime) -> typing.Mapping[str, str | object]:
my_data = config["data"]["personal-data"]
events += birthday.get_birthdays(last_year, os.path.join(my_data, "entities.yaml"))
events += get_accommodation(today, os.path.join(my_data, "accommodation.yaml"))
events += get_all_travel_events(today)
events += travel.all_events(today, config["data"]["personal-data"])
events += get_conferences(os.path.join(my_data, "conferences.yaml"))
events += waste_collection_events() + bristol_waste_collection_events(today)

57
agenda/travel.py Normal file
View file

@ -0,0 +1,57 @@
"""Travel."""
import os
import typing
from datetime import date
import yaml
from .types import Event
Leg = dict[str, str]
def get(
from_date: date,
method: str,
filepath: str,
extra: typing.Callable[[Leg], str] | None = None,
) -> list[Event]:
"""Get travel events."""
def title(item: Leg) -> str:
ret = f'{method} from {item["from"]} to {item["to"]}'
if extra:
ret += f" ({extra(item)})"
return ret
return [
Event(
date=item["depart"],
end_date=item["arrive"],
name="transport",
title=title(item),
url=item.get("url"),
)
for item in yaml.safe_load(open(filepath))
if item["depart"].date() >= from_date
]
def flight_number(flight: Leg) -> str:
"""Flight number."""
airline_code = flight["airline"]
# make sure this is the airline code, not the airline name
assert " " not in airline_code and not any(c.islower() for c in airline_code)
return airline_code + flight["flight_number"]
def all_events(from_date: date, data_dir: str) -> list[Event]:
"""Get all flights and rail journeys."""
trains = get(from_date, "train", os.path.join(data_dir, "trains.yaml"))
flights = get(
from_date, "flight", os.path.join(data_dir, "flights.yaml"), extra=flight_number
)
return trains + flights