Merge pull request #15 from warrior25/gtfs-service-day

GTFS service day handling

Author: warrior25

README.md

@@ -27,6 +27,7 @@ Each station creates a sensor which contains data for departures from that stati
 | last_refresh | Timestamp (ISO 8601 format) indicating when real-time departures were last fetched. |
 | departures   | A list of departure objects representing the next available departures.             |
 | station_name | Name of the monitored stop.                                                         |
+| station_id   | Unique identifier of the monitored stop.                                            | 
 
 ### Departures
 

custom_components/nysse/__init__.py

@@ -15,7 +15,7 @@ async def async_setup_entry(
     hass.data[DOMAIN][entry.entry_id] = entry.data
 
     # Forward the setup to the sensor platform.
-    await hass.config_entries.async_forward_entry_setup(entry, "sensor")
+    await hass.config_entries.async_forward_entry_setups(entry, ["sensor"])
     entry.async_on_unload(entry.add_update_listener(options_update_listener))
     return True
 

custom_components/nysse/fetch_api.py

@@ -7,9 +7,13 @@
 import os
 import pathlib
 import sqlite3
+from typing import NamedTuple
 import zipfile
 
 import aiohttp
+from dateutil import parser
+
+import homeassistant.util.dt as dt_util
 
 from .const import DOMAIN, GTFS_URL
 
@@ -240,8 +244,7 @@ def _get_file_modified_time(file_path):
 def _parse_csv_file(file_path):
     with open(file_path, newline="", encoding="utf-8") as csvfile:
         reader = csv.DictReader(csvfile)
-        data = [row.copy() for row in reader]
-    return data
+        return [row.copy() for row in reader]
 
 
 async def get_stops():
@@ -289,6 +292,16 @@ async def get_route_ids(stop_id):
     return route_ids
 
 
+class StopTime(NamedTuple):
+    route_id: str
+    trip_headsign: str
+    departure_time: datetime
+    aimed_departure_time: datetime | None
+    delay: int | None
+    delta_days: int
+    realtime: bool
+
+
 async def get_stop_times(stop_id, route_ids, amount, from_time):
     """Get the stop times for a given stop ID, route IDs, and amount.
 
@@ -306,26 +319,52 @@ async def get_stop_times(stop_id, route_ids, amount, from_time):
     conn, cursor = _get_database()
     today = datetime.now().strftime("%Y%m%d")
     weekday = datetime.strptime(today, "%Y%m%d").strftime("%A").lower()
-    stop_times = []
+    stop_times: list[StopTime] = []
     delta_days = 0
+    start_time = from_time.strftime("%H:%M:%S")
     while len(stop_times) < amount:
         cursor.execute(
             f"""
-            SELECT stop_times.trip_id, route_id, trip_headsign, departure_time, {delta_days} as delta_days
+            SELECT route_id, trip_headsign, departure_time
             FROM stop_times
             JOIN trips ON stop_times.trip_id = trips.trip_id
             JOIN calendar ON trips.service_id = calendar.service_id
             WHERE stop_id = ?
-            AND trips.route_id IN ({','.join(['?']*len(route_ids))})
+            AND trips.route_id IN ({",".join(["?"] * len(route_ids))})
             AND calendar.{weekday} = '1'
             AND calendar.start_date <= ?
             AND calendar.end_date >= ?
             AND departure_time > ?
             LIMIT ?
             """,
-            [stop_id, *route_ids, today, today, from_time.strftime("%H:%M:%S"), amount],
+            [stop_id, *route_ids, today, today, start_time, amount],
         )
-        stop_times += cursor.fetchall()
+
+        for row in cursor.fetchall():
+            route_id, trip_headsign, departure_time_str = row
+            row_delta_days = delta_days
+            hours, minutes, seconds = map(int, departure_time_str.split(":"))
+
+            if hours > 23:
+                hours -= 24
+                row_delta_days += 1
+
+            valid_time_str = f"{hours:02}:{minutes:02}:{seconds:02}"
+
+            departure_time = dt_util.as_local(parser.parse(valid_time_str))
+
+            stop_times.append(
+                StopTime(
+                    route_id,
+                    trip_headsign,
+                    departure_time,
+                    None,
+                    None,
+                    row_delta_days,
+                    False,
+                )
+            )
+
         if len(stop_times) >= amount:
             break
         # If there are no more stop times for today, move to the next day
@@ -338,5 +377,6 @@ async def get_stop_times(stop_id, route_ids, amount, from_time):
         next_day = datetime.strptime(today, "%Y%m%d") + timedelta(days=1)
         today = next_day.strftime("%Y%m%d")
         weekday = next_day.strftime("%A").lower()
+        start_time = "00:00:00"
     conn.close()
     return stop_times[:amount]

custom_components/nysse/sensor.py

@@ -7,6 +7,7 @@
 import logging
 
 from dateutil import parser
+from dateutil.parser import ParserError
 import isodate
 
 from homeassistant import config_entries, core
@@ -24,7 +25,7 @@
     STOP_URL,
     TRAM_LINES,
 )
-from .fetch_api import get_stop_times, get_stops
+from .fetch_api import StopTime, get_stop_times, get_stops
 from .network import get
 
 _LOGGER = logging.getLogger(__name__)
@@ -81,19 +82,17 @@ def __init__(self, stop_code, maximum, timelimit, lines) -> None:
 
         self._last_update_time = None
 
-    def _remove_unwanted_departures(self, departures):
+    def _remove_unwanted_departures(self, departures: list[StopTime]):
         try:
             removed_departures_count = 0
 
             # Remove unwanted departures based on departure time and line number
             for departure in departures[:]:
-                departure_local = dt_util.as_local(
-                    parser.parse(departure["departure_time"])
-                )
+                departure_local = dt_util.as_local(departure.departure_time)
                 if (
                     departure_local
                     < self._last_update_time + timedelta(minutes=self._timelimit)
-                    or departure["route_id"] not in self._lines
+                    or departure.route_id not in self._lines
                 ):
                     departures.remove(departure)
                     removed_departures_count += 1
@@ -129,7 +128,7 @@ async def _fetch_departures(self):
                     self._stop_code,
                     url,
                 )
-                return
+                return None
             unformatted_departures = json.loads(data)
             return self._format_departures(unformatted_departures)
         except OSError as err:
@@ -139,25 +138,20 @@ async def _fetch_departures(self):
     def _format_departures(self, departures):
         try:
             body = departures["body"][self._stop_code]
-            formatted_data = []
+            formatted_data: list[StopTime] = []
             for departure in body:
                 try:
-                    formatted_departure = {
-                        "route_id": departure["lineRef"],
-                        "trip_headsign": self._get_stop_name(
-                            departure["destinationShortName"]
-                        ),
-                        "departure_time": departure["call"]["expectedDepartureTime"],
-                        "aimed_departure_time": departure["call"]["aimedDepartureTime"],
-                        "delay": departure["delay"],
-                        "realtime": True,
-                    }
-                    if (
-                        formatted_departure["departure_time"] is not None
-                        and formatted_departure["aimed_departure_time"] is not None
-                    ):
-                        formatted_data.append(formatted_departure)
-                except KeyError as err:
+                    formatted_departure = StopTime(
+                        departure["lineRef"],
+                        self._get_stop_name(departure["destinationShortName"]),
+                        parser.parse(departure["call"]["expectedDepartureTime"]),
+                        parser.parse(departure["call"]["aimedDepartureTime"]),
+                        self._delay_to_display_format(departure["delay"]),
+                        0,
+                        True,
+                    )
+                    formatted_data.append(formatted_departure)
+                except (KeyError, ParserError) as err:
                     _LOGGER.info(
                         "%s: Failed to process realtime departure: %s",
                         self._stop_code,
@@ -200,13 +194,9 @@ async def async_update(self) -> None:
                 )
                 for journey in self._journeys[:]:
                     for departure in departures:
-                        departure_time = parser.parse(departure["aimed_departure_time"])
-                        journey_time = dt_util.as_local(
-                            parser.parse(journey["departure_time"])
-                        )
                         if (
-                            journey_time == departure_time
-                            and journey["route_id"] == departure["route_id"]
+                            journey.departure_time == departure.aimed_departure_time
+                            and journey.route_id == departure.route_id
                         ):
                             self._journeys.remove(journey)
             else:
@@ -223,24 +213,24 @@ async def async_update(self) -> None:
         except (OSError, ValueError) as err:
             _LOGGER.error("%s: Failed to update sensor: %s", self._stop_code, err)
 
-    def _data_to_display_format(self, data):
+    def _data_to_display_format(self, data: list[StopTime]):
         try:
             formatted_data = []
             for item in data:
                 departure = {
-                    "destination": item["trip_headsign"],
-                    "line": item["route_id"],
-                    "departure": parser.parse(item["departure_time"]).strftime("%H:%M"),
+                    "destination": item.trip_headsign,
+                    "line": item.route_id,
+                    "departure": item.departure_time.strftime("%H:%M"),
                     "time_to_station": self._time_to_station(item),
-                    "icon": self._get_line_icon(item["route_id"]),
-                    "realtime": item["realtime"] if "realtime" in item else False,
+                    "icon": self._get_line_icon(item.route_id),
+                    "realtime": item.realtime,
                 }
-                if "aimed_departure_time" in item:
-                    departure["aimed_departure"] = parser.parse(
-                        item["aimed_departure_time"]
-                    ).strftime("%H:%M")
-                if "delay" in item:
-                    departure["delay"] = self._delay_to_display_format(item["delay"])
+                if item.aimed_departure_time is not None:
+                    departure["aimed_departure"] = item.aimed_departure_time.strftime(
+                        "%H:%M"
+                    )
+                if item.delay is not None:
+                    departure["delay"] = item.delay
                 formatted_data.append(departure)
             return sorted(formatted_data, key=lambda x: x["time_to_station"])
         except (OSError, ValueError) as err:
@@ -252,11 +242,11 @@ def _get_line_icon(self, line_no):
             return "mdi:tram"
         return "mdi:bus"
 
-    def _time_to_station(self, item):
+    def _time_to_station(self, item: StopTime):
         try:
-            departure_local = dt_util.as_local(parser.parse(item["departure_time"]))
-            if "delta_days" in item:
-                departure_local += timedelta(days=item["delta_days"])
+            departure_local = dt_util.as_local(item.departure_time)
+            if item.delta_days > 0:
+                departure_local += timedelta(days=item.delta_days)
             next_departure_time = (departure_local - self._last_update_time).seconds
             return int(next_departure_time / 60)
         except OSError as err:
@@ -323,12 +313,12 @@ def state(self) -> str:
     @property
     def extra_state_attributes(self):
         """Sensor attributes."""
-        attributes = {
+        return {
             "last_refresh": self._last_update_time,
             "departures": self._all_data,
             "station_name": self._get_stop_name(self._stop_code),
+            "station_id": self._stop_code,
         }
-        return attributes
 
 
 class ServiceAlertSensor(SensorEntity):
@@ -363,7 +353,7 @@ async def _fetch_service_alerts(self):
                     "Nysse API error: failed to fetch service alerts: no data received from %s",
                     SERVICE_ALERTS_URL,
                 )
-                return
+                return None
             json_data = json.loads(data)
 
             self._last_update = self._timestamp_to_local(
@@ -428,8 +418,7 @@ def state(self) -> str:
     @property
     def extra_state_attributes(self):
         """Sensor attributes."""
-        attributes = {
+        return {
             "last_refresh": self._last_update,
             "alerts": self._alerts,
         }
-        return attributes