diff --git a/emt_madrid/domain/line.py b/emt_madrid/domain/line.py index 946adfb..c7342e1 100644 --- a/emt_madrid/domain/line.py +++ b/emt_madrid/domain/line.py @@ -18,6 +18,7 @@ class Line: day_type: Optional[DayType] = None arrival: Optional[int] = None next_arrival: Optional[int] = None + bus_coords: Optional[list[float]] = None def __str__(self) -> str: """Return a string representation of the line.""" diff --git a/emt_madrid/infrastructure/emt_api_repository.py b/emt_madrid/infrastructure/emt_api_repository.py index d1428b9..8434f59 100644 --- a/emt_madrid/infrastructure/emt_api_repository.py +++ b/emt_madrid/infrastructure/emt_api_repository.py @@ -233,12 +233,22 @@ async def get_arrivals(self, stop: Stop) -> Stop: if not sorted_arrivals: line.arrival = None line.next_arrival = None + line.bus_coords = None else: line.arrival = sorted_arrivals[0] line.next_arrival = ( sorted_arrivals[1] if len(sorted_arrivals) > 1 else None ) + bus_for_line = [ + a for a in arrivals_data if str(a["line"]) == line.line_number + ] + if bus_for_line: + first_bus = min(bus_for_line, key=lambda x: int(x["estimateArrive"])) + line.bus_coords = first_bus["geometry"]["coordinates"] + else: + line.bus_coords = None + return stop except Exception as e: