Refactor
This commit is contained in:
parent
135476b59f
commit
bdc94f3ebc
458
agenda/busy.py
458
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"))
|
||||
|
||||
|
|
|
|||
Loading…
Reference in a new issue