Refactor and add some docstrings.

This commit is contained in:
Edward Betts 2024-10-02 10:16:30 +01:00
parent 8cbfb745c4
commit 0b23f71aa6

View file

@ -136,42 +136,29 @@ def load_flights(flight_bookings: list[StrDict]) -> list[StrDict]:
return flights return flights
def build_trip_list( def collect_travel_items(
flight_bookings: list[StrDict],
data_dir: str | None = None, data_dir: str | None = None,
route_distances: travel.RouteDistances | None = None, route_distances: travel.RouteDistances | None = None,
) -> list[Trip]: ) -> list[StrDict]:
"""Generate list of trips.""" """Generate list of trips."""
trips: dict[date, Trip] = {}
if data_dir is None: if data_dir is None:
data_dir = flask.current_app.config["PERSONAL_DATA"] data_dir = flask.current_app.config["PERSONAL_DATA"]
yaml_trip_list = travel.parse_yaml("trips", data_dir) return sorted(
load_flights(load_flight_bookings(data_dir))
yaml_trip_lookup = {item["trip"]: item for item in yaml_trip_list}
flight_bookings = load_flight_bookings(data_dir)
travel_items = sorted(
load_flights(flight_bookings)
+ load_trains(data_dir, route_distances=route_distances) + load_trains(data_dir, route_distances=route_distances)
+ load_ferries(data_dir, route_distances=route_distances), + load_ferries(data_dir, route_distances=route_distances),
key=depart_datetime, key=depart_datetime,
) )
data = {
"flight_bookings": flight_bookings,
"travel": travel_items,
"accommodation": travel.parse_yaml("accommodation", data_dir),
"conferences": travel.parse_yaml("conferences", data_dir),
"events": travel.parse_yaml("events", data_dir),
}
for item in data["accommodation"]:
price = item.get("price")
if price:
item["price"] = decimal.Decimal(price)
def group_travel_items_into_trips(
data: StrDict, yaml_trip_list: list[StrDict]
) -> list[Trip]:
"""Group travel items into trips."""
trips: dict[date, Trip] = {}
yaml_trip_lookup = {item["trip"]: item for item in yaml_trip_list}
for key, item_list in data.items(): for key, item_list in data.items():
assert isinstance(item_list, list) assert isinstance(item_list, list)
for item in item_list: for item in item_list:
@ -187,6 +174,34 @@ def build_trip_list(
return [trip for _, trip in sorted(trips.items())] return [trip for _, trip in sorted(trips.items())]
def build_trip_list(
data_dir: str | None = None,
route_distances: travel.RouteDistances | None = None,
) -> list[Trip]:
"""Generate list of trips."""
if data_dir is None:
data_dir = flask.current_app.config["PERSONAL_DATA"]
yaml_trip_list = travel.parse_yaml("trips", data_dir)
flight_bookings = load_flight_bookings(data_dir)
data = {
"flight_bookings": flight_bookings,
"travel": collect_travel_items(flight_bookings, data_dir, route_distances),
"accommodation": travel.parse_yaml("accommodation", data_dir),
"conferences": travel.parse_yaml("conferences", data_dir),
"events": travel.parse_yaml("events", data_dir),
}
for item in data["accommodation"]:
price = item.get("price")
if price:
item["price"] = decimal.Decimal(price)
return group_travel_items_into_trips(data, yaml_trip_list)
def add_coordinates_for_unbooked_flights( def add_coordinates_for_unbooked_flights(
routes: list[StrDict], coordinates: list[StrDict] routes: list[StrDict], coordinates: list[StrDict]
) -> None: ) -> None:
@ -211,7 +226,7 @@ def add_coordinates_for_unbooked_flights(
def collect_trip_coordinates(trip: Trip) -> list[StrDict]: def collect_trip_coordinates(trip: Trip) -> list[StrDict]:
"""Extract and deduplicate airport and station coordinates from trip.""" """Extract and de-duplicate airport and station coordinates from trip."""
stations = {} stations = {}
station_list = [] station_list = []
airports = {} airports = {}
@ -368,6 +383,7 @@ def get_trip_routes(trip: Trip) -> list[StrDict]:
def get_coordinates_and_routes( def get_coordinates_and_routes(
trip_list: list[Trip], data_dir: str | None = None trip_list: list[Trip], data_dir: str | None = None
) -> tuple[list[StrDict], list[StrDict]]: ) -> tuple[list[StrDict], list[StrDict]]:
"""Given a list of trips return the associated coordinates and routes."""
if data_dir is None: if data_dir is None:
data_dir = flask.current_app.config["PERSONAL_DATA"] data_dir = flask.current_app.config["PERSONAL_DATA"]
coordinates = [] coordinates = []