Add train and ferry support to location tracking
- Add helper functions for train and ferry location extraction - Update get_location_for_date to consider trains and ferries alongside flights - Parse stations.yaml and ferry_terminals.yaml for location data - Handle UK vs international locations consistently for all transport modes - Add comprehensive tests for new train and ferry location helpers - Format code with black for consistent style Now tracks complete travel history including flights, trains, ferries, and accommodation for accurate location determination. 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
parent
f89d984623
commit
e370049bcb
246
agenda/busy.py
246
agenda/busy.py
|
|
@ -151,12 +151,79 @@ def _get_accommodation_location(
|
|||
return (acc.get("location", "Unknown"), get_country(acc.get("country", "gb")))
|
||||
|
||||
|
||||
def _get_train_location(
|
||||
train_leg: StrDict, stations: StrDict, on_trip: bool = False
|
||||
) -> tuple[str | None, pycountry.db.Country | None]:
|
||||
"""Get location from train leg data."""
|
||||
destination = train_leg.get("to")
|
||||
if not destination:
|
||||
return (None, get_country("gb"))
|
||||
|
||||
# Find station info
|
||||
station_info = None
|
||||
for station in stations:
|
||||
if station.get("name") == destination:
|
||||
station_info = station
|
||||
break
|
||||
|
||||
if not station_info:
|
||||
return (destination, get_country("gb"))
|
||||
|
||||
station_country = station_info.get("country", "gb")
|
||||
|
||||
if station_country == "gb":
|
||||
if on_trip:
|
||||
# When on a trip, show the actual location even for UK stations
|
||||
return (destination, get_country("gb"))
|
||||
else:
|
||||
# When not on a trip, UK stations mean home
|
||||
return (None, get_country("gb"))
|
||||
else:
|
||||
return (destination, get_country(station_country))
|
||||
|
||||
|
||||
def _get_ferry_location(
|
||||
ferry: StrDict, terminals: StrDict, on_trip: bool = False
|
||||
) -> tuple[str | None, pycountry.db.Country | None]:
|
||||
"""Get location from ferry data."""
|
||||
destination = ferry.get("to")
|
||||
if not destination:
|
||||
return (None, get_country("gb"))
|
||||
|
||||
# Find terminal info
|
||||
terminal_info = None
|
||||
for terminal in terminals:
|
||||
if terminal.get("name") == destination:
|
||||
terminal_info = terminal
|
||||
break
|
||||
|
||||
if not terminal_info:
|
||||
return (destination, get_country("gb"))
|
||||
|
||||
terminal_country = terminal_info.get("country", "gb")
|
||||
terminal_city = terminal_info.get("city", destination)
|
||||
|
||||
if terminal_country == "gb":
|
||||
if on_trip:
|
||||
# When on a trip, show the actual location even for UK terminals
|
||||
return (terminal_city, get_country("gb"))
|
||||
else:
|
||||
# When not on a trip, UK terminals mean home
|
||||
return (None, get_country("gb"))
|
||||
else:
|
||||
return (terminal_city, get_country(terminal_country))
|
||||
|
||||
|
||||
def _find_most_recent_travel_within_trip(
|
||||
trip: Trip,
|
||||
target_date: date,
|
||||
bookings: list[StrDict],
|
||||
accommodations: list[StrDict],
|
||||
airports: StrDict,
|
||||
trains: list[StrDict] | None = None,
|
||||
stations: StrDict | None = None,
|
||||
ferries: list[StrDict] | None = None,
|
||||
terminals: StrDict | None = None,
|
||||
) -> tuple[str | None, pycountry.db.Country | None] | None:
|
||||
"""Find the most recent travel location within a trip."""
|
||||
uk_airports = {"LHR", "LGW", "STN", "LTN", "BRS", "BHX", "MAN", "EDI", "GLA"}
|
||||
|
|
@ -219,6 +286,69 @@ def _find_most_recent_travel_within_trip(
|
|||
acc, on_trip=True
|
||||
)
|
||||
|
||||
# Check trains within trip period
|
||||
if trains and stations:
|
||||
for train in trains:
|
||||
for leg in train.get("legs", []):
|
||||
if "arrive" in leg:
|
||||
try:
|
||||
arrive_datetime, arrive_date = _parse_datetime_field(
|
||||
leg["arrive"]
|
||||
)
|
||||
except ValueError:
|
||||
continue
|
||||
|
||||
# Only consider trains within this trip and before target date
|
||||
if trip.start <= arrive_date <= target_date:
|
||||
# Compare both date and time to handle same-day arrivals correctly
|
||||
if (
|
||||
trip_most_recent_date is None
|
||||
or arrive_date > trip_most_recent_date
|
||||
or (
|
||||
arrive_date == trip_most_recent_date
|
||||
and (
|
||||
trip_most_recent_datetime is None
|
||||
or arrive_datetime > trip_most_recent_datetime
|
||||
)
|
||||
)
|
||||
):
|
||||
trip_most_recent_date = arrive_date
|
||||
trip_most_recent_datetime = arrive_datetime
|
||||
trip_most_recent_location = _get_train_location(
|
||||
leg, stations, on_trip=True
|
||||
)
|
||||
|
||||
# Check ferries within trip period
|
||||
if ferries and terminals:
|
||||
for ferry in ferries:
|
||||
if "arrive" in ferry:
|
||||
try:
|
||||
arrive_datetime, arrive_date = _parse_datetime_field(
|
||||
ferry["arrive"]
|
||||
)
|
||||
except ValueError:
|
||||
continue
|
||||
|
||||
# Only consider ferries within this trip and before target date
|
||||
if trip.start <= arrive_date <= target_date:
|
||||
# Compare both date and time to handle same-day arrivals correctly
|
||||
if (
|
||||
trip_most_recent_date is None
|
||||
or arrive_date > trip_most_recent_date
|
||||
or (
|
||||
arrive_date == trip_most_recent_date
|
||||
and (
|
||||
trip_most_recent_datetime is None
|
||||
or arrive_datetime > trip_most_recent_datetime
|
||||
)
|
||||
)
|
||||
):
|
||||
trip_most_recent_date = arrive_date
|
||||
trip_most_recent_datetime = arrive_datetime
|
||||
trip_most_recent_location = _get_ferry_location(
|
||||
ferry, terminals, on_trip=True
|
||||
)
|
||||
|
||||
return trip_most_recent_location
|
||||
|
||||
|
||||
|
|
@ -253,6 +383,10 @@ def _find_most_recent_travel_before_date(
|
|||
bookings: list[StrDict],
|
||||
accommodations: list[StrDict],
|
||||
airports: StrDict,
|
||||
trains: list[StrDict] | None = None,
|
||||
stations: StrDict | None = None,
|
||||
ferries: list[StrDict] | None = None,
|
||||
terminals: StrDict | None = None,
|
||||
) -> tuple[str | None, pycountry.db.Country | None] | None:
|
||||
"""Find the most recent travel location before a given date."""
|
||||
uk_airports = {"LHR", "LGW", "STN", "LTN", "BRS", "BHX", "MAN", "EDI", "GLA"}
|
||||
|
|
@ -308,6 +442,67 @@ def _find_most_recent_travel_before_date(
|
|||
acc, on_trip=False
|
||||
)
|
||||
|
||||
# Check trains
|
||||
if trains and stations:
|
||||
for train in trains:
|
||||
for leg in train.get("legs", []):
|
||||
if "arrive" in leg:
|
||||
try:
|
||||
arrive_datetime, arrive_date = _parse_datetime_field(
|
||||
leg["arrive"]
|
||||
)
|
||||
except ValueError:
|
||||
continue
|
||||
|
||||
if arrive_date <= target_date:
|
||||
# Compare both date and time to handle same-day arrivals correctly
|
||||
if (
|
||||
most_recent_date is None
|
||||
or arrive_date > most_recent_date
|
||||
or (
|
||||
arrive_date == most_recent_date
|
||||
and (
|
||||
most_recent_datetime is None
|
||||
or arrive_datetime > most_recent_datetime
|
||||
)
|
||||
)
|
||||
):
|
||||
most_recent_date = arrive_date
|
||||
most_recent_datetime = arrive_datetime
|
||||
most_recent_location = _get_train_location(
|
||||
leg, stations, on_trip=False
|
||||
)
|
||||
|
||||
# Check ferries
|
||||
if ferries and terminals:
|
||||
for ferry in ferries:
|
||||
if "arrive" in ferry:
|
||||
try:
|
||||
arrive_datetime, arrive_date = _parse_datetime_field(
|
||||
ferry["arrive"]
|
||||
)
|
||||
except ValueError:
|
||||
continue
|
||||
|
||||
if arrive_date <= target_date:
|
||||
# Compare both date and time to handle same-day arrivals correctly
|
||||
if (
|
||||
most_recent_date is None
|
||||
or arrive_date > most_recent_date
|
||||
or (
|
||||
arrive_date == most_recent_date
|
||||
and (
|
||||
most_recent_datetime is None
|
||||
or arrive_datetime > most_recent_datetime
|
||||
)
|
||||
)
|
||||
):
|
||||
most_recent_date = arrive_date
|
||||
most_recent_datetime = arrive_datetime
|
||||
most_recent_location = _get_ferry_location(
|
||||
ferry, terminals, on_trip=False
|
||||
)
|
||||
|
||||
return most_recent_location
|
||||
|
||||
|
||||
|
|
@ -349,14 +544,26 @@ def get_location_for_date(
|
|||
bookings: list[StrDict],
|
||||
accommodations: list[StrDict],
|
||||
airports: StrDict,
|
||||
trains: list[StrDict] | None = None,
|
||||
stations: StrDict | None = None,
|
||||
ferries: list[StrDict] | None = None,
|
||||
terminals: StrDict | None = None,
|
||||
) -> tuple[str | None, pycountry.db.Country | None]:
|
||||
"""Get location (city, country) for a specific date using travel history."""
|
||||
# First check if currently on a trip
|
||||
for trip in trips:
|
||||
if trip.start <= target_date <= (trip.end or trip.start):
|
||||
# For trips, find the most recent flight or accommodation within the trip period
|
||||
# For trips, find the most recent travel within the trip period
|
||||
trip_location = _find_most_recent_travel_within_trip(
|
||||
trip, target_date, bookings, accommodations, airports
|
||||
trip,
|
||||
target_date,
|
||||
bookings,
|
||||
accommodations,
|
||||
airports,
|
||||
trains,
|
||||
stations,
|
||||
ferries,
|
||||
terminals,
|
||||
)
|
||||
if trip_location:
|
||||
return trip_location
|
||||
|
|
@ -366,9 +573,16 @@ def get_location_for_date(
|
|||
if progression_location:
|
||||
return progression_location
|
||||
|
||||
# Find most recent flight or accommodation before this date
|
||||
# Find most recent travel before this date
|
||||
recent_travel = _find_most_recent_travel_before_date(
|
||||
target_date, bookings, accommodations, airports
|
||||
target_date,
|
||||
bookings,
|
||||
accommodations,
|
||||
airports,
|
||||
trains,
|
||||
stations,
|
||||
ferries,
|
||||
terminals,
|
||||
)
|
||||
|
||||
# Check for recent trips that have ended - prioritize this over individual travel data
|
||||
|
|
@ -400,6 +614,10 @@ def weekends(
|
|||
bookings = travel.parse_yaml("flights", data_dir)
|
||||
accommodations = travel.parse_yaml("accommodation", data_dir)
|
||||
airports = travel.parse_yaml("airports", data_dir)
|
||||
trains = travel.parse_yaml("trains", data_dir)
|
||||
stations = travel.parse_yaml("stations", data_dir)
|
||||
ferries = travel.parse_yaml("ferries", data_dir)
|
||||
terminals = travel.parse_yaml("ferry_terminals", data_dir)
|
||||
|
||||
weekends_info = []
|
||||
for i in range(52):
|
||||
|
|
@ -418,10 +636,26 @@ def weekends(
|
|||
]
|
||||
|
||||
saturday_location = get_location_for_date(
|
||||
saturday, trips, bookings, accommodations, airports
|
||||
saturday,
|
||||
trips,
|
||||
bookings,
|
||||
accommodations,
|
||||
airports,
|
||||
trains,
|
||||
stations,
|
||||
ferries,
|
||||
terminals,
|
||||
)
|
||||
sunday_location = get_location_for_date(
|
||||
sunday, trips, bookings, accommodations, airports
|
||||
sunday,
|
||||
trips,
|
||||
bookings,
|
||||
accommodations,
|
||||
airports,
|
||||
trains,
|
||||
stations,
|
||||
ferries,
|
||||
terminals,
|
||||
)
|
||||
|
||||
weekends_info.append(
|
||||
|
|
|
|||
|
|
@ -178,6 +178,68 @@ def test_parse_datetime_field():
|
|||
assert parsed_dt.day == 1
|
||||
|
||||
|
||||
def test_train_location_helpers():
|
||||
"""Test the train location helper functions."""
|
||||
from agenda.busy import _get_train_location
|
||||
|
||||
# Mock station data
|
||||
stations = [
|
||||
{"name": "London St Pancras", "country": "gb"},
|
||||
{"name": "Brussels Midi", "country": "be"},
|
||||
{"name": "Edinburgh Waverley", "country": "gb"},
|
||||
]
|
||||
|
||||
# Test UK station when not on trip (should return None for home)
|
||||
train_leg = {"to": "London St Pancras"}
|
||||
location = _get_train_location(train_leg, stations, on_trip=False)
|
||||
assert location[0] is None # Should be home
|
||||
assert location[1].alpha_2 == "GB"
|
||||
|
||||
# Test UK station when on trip (should return city name)
|
||||
location = _get_train_location(train_leg, stations, on_trip=True)
|
||||
assert location[0] == "London St Pancras"
|
||||
assert location[1].alpha_2 == "GB"
|
||||
|
||||
# Test non-UK station
|
||||
train_leg = {"to": "Brussels Midi"}
|
||||
location = _get_train_location(train_leg, stations, on_trip=False)
|
||||
assert location[0] == "Brussels Midi"
|
||||
assert location[1].alpha_2 == "BE"
|
||||
|
||||
|
||||
def test_ferry_location_helpers():
|
||||
"""Test the ferry location helper functions."""
|
||||
from agenda.busy import _get_ferry_location
|
||||
|
||||
# Mock terminal data
|
||||
terminals = [
|
||||
{"name": "Dover Eastern Docks", "country": "gb", "city": "Dover"},
|
||||
{"name": "Calais Ferry Terminal", "country": "fr", "city": "Calais"},
|
||||
{
|
||||
"name": "Portsmouth Continental Terminal",
|
||||
"country": "gb",
|
||||
"city": "Portsmouth",
|
||||
},
|
||||
]
|
||||
|
||||
# Test UK terminal when not on trip (should return None for home)
|
||||
ferry = {"to": "Dover Eastern Docks"}
|
||||
location = _get_ferry_location(ferry, terminals, on_trip=False)
|
||||
assert location[0] is None # Should be home
|
||||
assert location[1].alpha_2 == "GB"
|
||||
|
||||
# Test UK terminal when on trip (should return city name)
|
||||
location = _get_ferry_location(ferry, terminals, on_trip=True)
|
||||
assert location[0] == "Dover"
|
||||
assert location[1].alpha_2 == "GB"
|
||||
|
||||
# Test non-UK terminal
|
||||
ferry = {"to": "Calais Ferry Terminal"}
|
||||
location = _get_ferry_location(ferry, terminals, on_trip=False)
|
||||
assert location[0] == "Calais"
|
||||
assert location[1].alpha_2 == "FR"
|
||||
|
||||
|
||||
def test_get_busy_events(app_context, trips):
|
||||
"""Test get_busy_events function."""
|
||||
start_date = date(2023, 1, 1)
|
||||
|
|
|
|||
Loading…
Reference in a new issue