From e370049bcb665ad80b1d088da0e132a6e2f035ad Mon Sep 17 00:00:00 2001 From: Edward Betts Date: Wed, 16 Jul 2025 11:10:58 +0200 Subject: [PATCH] Add train and ferry support to location tracking MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 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 --- agenda/busy.py | 246 +++++++++++++++++++++++++++++++++++++++++++-- tests/test_busy.py | 62 ++++++++++++ 2 files changed, 302 insertions(+), 6 deletions(-) diff --git a/agenda/busy.py b/agenda/busy.py index 4fcbeec..31cc4b6 100644 --- a/agenda/busy.py +++ b/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( diff --git a/tests/test_busy.py b/tests/test_busy.py index 150dc8f..026eeb9 100644 --- a/tests/test_busy.py +++ b/tests/test_busy.py @@ -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)