diff --git a/agenda/busy.py b/agenda/busy.py index 87ccf1d..4fcbeec 100644 --- a/agenda/busy.py +++ b/agenda/busy.py @@ -11,6 +11,20 @@ from . import events_yaml, get_country, travel from .event import Event from .types import StrDict, Trip +NEARBY_BALKAN_COUNTRIES = { + "GR", + "AL", + "XK", + "HR", + "SI", + "MK", + "BA", + "ME", + "RS", + "BG", + "RO", +} + def busy_event(e: Event) -> bool: """Busy.""" @@ -81,167 +95,182 @@ def get_busy_events( return busy_events -def get_location_for_date( +def _parse_datetime_field(datetime_obj: datetime | date) -> tuple[datetime, date]: + """Parse a datetime field that could be datetime object or string.""" + if hasattr(datetime_obj, "date"): + return datetime_obj, datetime_obj.date() + elif isinstance(datetime_obj, str): + dt = datetime.fromisoformat(datetime_obj.replace("Z", "+00:00")) + return dt, dt.date() + else: + raise ValueError(f"Invalid datetime format: {datetime_obj}") + + +def _get_airport_location( + airport_code: str, airports: StrDict, uk_airports: set[str], on_trip: bool = False +) -> tuple[str | None, pycountry.db.Country | None]: + """Get location from airport code.""" + if airport_code in uk_airports: + if on_trip: + # When on a trip, show the actual location even for UK airports + airport_info = airports.get(airport_code) + if airport_info: + location_name = airport_info.get( + "city", airport_info.get("name", "London") + ) + return (location_name, get_country("gb")) + else: + return ("London", get_country("gb")) + else: + # When not on a trip, UK airports mean home + return (None, get_country("gb")) + else: + # Non-UK airports + airport_info = airports.get(airport_code) + if airport_info: + location_name = airport_info.get( + "city", airport_info.get("name", airport_code) + ) + return (location_name, get_country(airport_info.get("country", "gb"))) + else: + return (airport_code, get_country("gb")) + + +def _get_accommodation_location( + acc: StrDict, on_trip: bool = False +) -> tuple[str | None, pycountry.db.Country | None]: + """Get location from accommodation data.""" + if acc.get("country") == "gb": + if on_trip: + # When on a trip, show the actual location even for UK accommodations + return (acc.get("location", "London"), get_country("gb")) + else: + # When not on a trip, UK accommodation means home + return (None, get_country("gb")) + else: + return (acc.get("location", "Unknown"), get_country(acc.get("country", "gb"))) + + +def _find_most_recent_travel_within_trip( + trip: Trip, target_date: date, - trips: list[Trip], bookings: list[StrDict], accommodations: list[StrDict], airports: StrDict, -) -> tuple[str | None, pycountry.db.Country | None]: - """Get location (city, country) for a specific date using travel history.""" - # UK airports that indicate being home +) -> 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"} - # 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 - # to determine exact location on the target date - trip_most_recent_date = None - trip_most_recent_location = None + trip_most_recent_date = None + trip_most_recent_location = None + trip_most_recent_datetime = None - # Check flights within trip period - trip_most_recent_datetime = None - for booking in bookings: - for flight in booking.get("flights", []): - if "arrive" in flight: - arrive_datetime_obj = flight["arrive"] - if hasattr(arrive_datetime_obj, "date"): - arrive_datetime = arrive_datetime_obj - arrive_date = arrive_datetime_obj.date() - elif isinstance(arrive_datetime_obj, str): - arrive_datetime = datetime.fromisoformat( - arrive_datetime_obj.replace("Z", "+00:00") - ) - arrive_date = arrive_datetime.date() - - # Only consider flights within this trip and before target date - if trip.start <= arrive_date <= target_date: - # Compare both date and time to handle same-day flights 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 - destination_airport = flight["to"] - - if destination_airport in uk_airports: - # When on a trip, show the actual location even for UK airports - airport_info = airports.get(destination_airport) - if airport_info: - location_name = airport_info.get( - "city", - airport_info.get("name", "London"), - ) - trip_most_recent_location = ( - location_name, - get_country("gb"), - ) - else: - trip_most_recent_location = ( - "London", - get_country("gb"), - ) - else: - airport_info = airports.get(destination_airport) - if airport_info: - location_name = airport_info.get( - "city", - airport_info.get( - "name", destination_airport - ), - ) - trip_most_recent_location = ( - location_name, - get_country( - airport_info.get("country", "gb") - ), - ) - - # Check accommodations within trip period - for acc in accommodations: - if "from" in acc: - acc_date = acc["from"] - if hasattr(acc_date, "date"): - acc_date = acc_date.date() - elif isinstance(acc_date, str): - acc_date = datetime.fromisoformat( - acc_date.replace("Z", "+00:00") - ).date() - - # Only consider accommodations within this trip and before/on target date - if trip.start <= acc_date <= target_date: - # Accommodation takes precedence over flights on the same date - # or if it's genuinely more recent - if ( - trip_most_recent_date is None - or acc_date > trip_most_recent_date - or acc_date == trip_most_recent_date - ): - trip_most_recent_date = acc_date - if acc.get("country") == "gb": - # When on a trip, show the actual location even for UK accommodations - trip_most_recent_location = ( - acc.get("location", "London"), - get_country("gb"), - ) - else: - trip_most_recent_location = ( - acc.get("location", "Unknown"), - get_country(acc.get("country", "gb")), - ) - - # Return the most recent location within the trip, or fallback to trip location by date - if trip_most_recent_location: - return trip_most_recent_location - - # Fallback: determine location based on trip progression and date - locations = trip.locations() - if locations: - # If only one location, use it (when on a trip, always show the location) - if len(locations) == 1: - city, country = locations[0] - return (city, country) - - # Multiple locations: use progression through the trip - trip_duration = (trip.end - trip.start).days + 1 - days_into_trip = (target_date - trip.start).days - - # Simple progression: first half at first location, second half at last location - if days_into_trip <= trip_duration // 2: - city, country = locations[0] - else: - city, country = locations[-1] - - return (city, country) - - # Find most recent flight or accommodation before this date - most_recent_location = None - most_recent_date = None - - # Check flights - most_recent_datetime = None + # Check flights within trip period for booking in bookings: for flight in booking.get("flights", []): if "arrive" in flight: - arrive_datetime_obj = flight["arrive"] - if hasattr(arrive_datetime_obj, "date"): - arrive_datetime = arrive_datetime_obj - arrive_date = arrive_datetime_obj.date() - elif isinstance(arrive_datetime_obj, str): - arrive_datetime = datetime.fromisoformat( - arrive_datetime_obj.replace("Z", "+00:00") + try: + arrive_datetime, arrive_date = _parse_datetime_field( + flight["arrive"] ) - arrive_date = arrive_datetime.date() + except ValueError: + continue + + # Only consider flights within this trip and before target date + if trip.start <= arrive_date <= target_date: + # Compare both date and time to handle same-day flights 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 + destination_airport = flight["to"] + trip_most_recent_location = _get_airport_location( + destination_airport, airports, uk_airports, on_trip=True + ) + + # Check accommodations within trip period + for acc in accommodations: + if "from" in acc: + try: + _, acc_date = _parse_datetime_field(acc["from"]) + except ValueError: + continue + + # Only consider accommodations within this trip and before/on target date + if trip.start <= acc_date <= target_date: + # Accommodation takes precedence over flights on the same date + # or if it's genuinely more recent + if ( + trip_most_recent_date is None + or acc_date > trip_most_recent_date + or acc_date == trip_most_recent_date + ): + trip_most_recent_date = acc_date + trip_most_recent_location = _get_accommodation_location( + acc, on_trip=True + ) + + return trip_most_recent_location + + +def _get_trip_location_by_progression( + trip: Trip, target_date: date +) -> tuple[str | None, pycountry.db.Country | None] | None: + """Determine location based on trip progression and date.""" + locations = trip.locations() + if not locations: + return None + + # If only one location, use it (when on a trip, always show the location) + if len(locations) == 1: + city, country = locations[0] + return (city, country) + + # Multiple locations: use progression through the trip + trip_duration = (trip.end - trip.start).days + 1 + days_into_trip = (target_date - trip.start).days + + # Simple progression: first half at first location, second half at last location + if days_into_trip <= trip_duration // 2: + city, country = locations[0] + else: + city, country = locations[-1] + + return (city, country) + + +def _find_most_recent_travel_before_date( + target_date: date, + bookings: list[StrDict], + accommodations: list[StrDict], + airports: StrDict, +) -> 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"} + + most_recent_location = None + most_recent_date = None + most_recent_datetime = None + + # Check flights + for booking in bookings: + for flight in booking.get("flights", []): + if "arrive" in flight: + try: + arrive_datetime, arrive_date = _parse_datetime_field( + flight["arrive"] + ) + except ValueError: + continue if arrive_date <= target_date: # Compare both date and time to handle same-day flights correctly @@ -256,58 +285,42 @@ def get_location_for_date( ) ) ): - most_recent_date = arrive_date most_recent_datetime = arrive_datetime destination_airport = flight["to"] - - # If arriving at UK airport, assume back home in Bristol - if destination_airport in uk_airports: - most_recent_location = (None, get_country("gb")) - else: - # Get destination airport location for non-UK arrivals - airport_info = airports.get(destination_airport) - if airport_info: - location_name = airport_info.get( - "city", - airport_info.get("name", destination_airport), - ) - most_recent_location = ( - location_name, - get_country(airport_info.get("country", "gb")), - ) + most_recent_location = _get_airport_location( + destination_airport, airports, uk_airports, on_trip=False + ) # Check accommodation - only override if accommodation is more recent for acc in accommodations: if "from" in acc: - acc_date = acc["from"] - if hasattr(acc_date, "date"): - acc_date = acc_date.date() - elif isinstance(acc_date, str): - acc_date = datetime.fromisoformat( - acc_date.replace("Z", "+00:00") - ).date() + try: + _, acc_date = _parse_datetime_field(acc["from"]) + except ValueError: + continue if acc_date <= target_date: # Only update if this accommodation is more recent than existing result if most_recent_date is None or acc_date > most_recent_date: most_recent_date = acc_date - # For UK accommodation, use Bristol as location - if acc.get("country") == "gb": - most_recent_location = (None, get_country("gb")) - else: - most_recent_location = ( - acc.get("location", "Unknown"), - get_country(acc.get("country", "gb")), - ) + most_recent_location = _get_accommodation_location( + acc, on_trip=False + ) - # Check for recent trips that have ended - prioritize this over individual travel data - # This handles cases where you're traveling home after a trip (e.g. stopovers, connections) + return most_recent_location + + +def _check_return_home_heuristic( + target_date: date, trips: list[Trip] +) -> tuple[str | None, pycountry.db.Country | None] | None: + """Check if should return home based on recent trips that have ended.""" for trip in trips: if trip.end and trip.end < target_date: locations = trip.locations() if locations: final_city, final_country = locations[-1] + final_alpha_2 = final_country.alpha_2 days_since_trip = (target_date - trip.end).days # If trip ended in UK, you should be home now @@ -316,38 +329,57 @@ def get_location_for_date( # For short trips to nearby countries or international trips # (ended >=1 day ago), assume returned home if no subsequent travel data - if ( - days_since_trip >= 1 - and hasattr(final_country, "alpha_2") - and ( - # European countries (close by rail/ferry) - final_country.alpha_2 - in {"BE", "NL", "FR", "DE", "CH", "AT", "IT", "ES"} - # Nearby Balkan countries - or final_country.alpha_2 - in { - "GR", - "AL", - "XK", - "HR", - "SI", - "MK", - "BA", - "ME", - "RS", - "BG", - "RO", - } - # International trips (assume return home after trip ends) - or final_country.alpha_2 - in {"US", "CA", "IN", "JP", "CN", "AU", "NZ", "BR", "AR", "ZA"} - ) + if days_since_trip >= 1 and ( + # European countries (close by rail/ferry) + final_alpha_2 in {"BE", "NL", "FR", "DE", "CH", "AT", "IT", "ES"} + # Nearby Balkan countries + or final_alpha_2 in NEARBY_BALKAN_COUNTRIES + # International trips (assume return home after trip ends) + or final_alpha_2 + in {"US", "CA", "IN", "JP", "CN", "AU", "NZ", "BR", "AR", "ZA"} ): return (None, get_country("gb")) - # Return most recent location or default to Bristol - if most_recent_location: - return most_recent_location + return None + + +def get_location_for_date( + target_date: date, + trips: list[Trip], + bookings: list[StrDict], + accommodations: list[StrDict], + airports: StrDict, +) -> 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 + trip_location = _find_most_recent_travel_within_trip( + trip, target_date, bookings, accommodations, airports + ) + if trip_location: + return trip_location + + # Fallback: determine location based on trip progression and date + progression_location = _get_trip_location_by_progression(trip, target_date) + if progression_location: + return progression_location + + # Find most recent flight or accommodation before this date + recent_travel = _find_most_recent_travel_before_date( + target_date, bookings, accommodations, airports + ) + + # Check for recent trips that have ended - prioritize this over individual travel data + # This handles cases where you're traveling home after a trip (e.g. stopovers, connections) + return_home = _check_return_home_heuristic(target_date, trips) + if return_home: + return return_home + + # Return most recent location or default to home + if recent_travel: + return recent_travel return (None, get_country("gb"))