Improve type hints

Author: ViViDboarder

py_nextbus/client.py

@@ -10,15 +10,21 @@
 import requests
 from requests.exceptions import HTTPError
 
+from py_nextbus.models import AgencyInfo
+from py_nextbus.models import RouteDetails
+from py_nextbus.models import RouteInfo
+from py_nextbus.models import StopPrediction
+
 LOG = logging.getLogger()
 
 
 class NextBusError(Exception):
     pass
 
 
-class NextBusHTTPError(HTTPError, NextBusError):
+class NextBusHTTPError(NextBusError):
     def __init__(self, message: str, http_err: HTTPError):
+        super().__init__()
         self.__dict__.update(http_err.__dict__)
         self.message: str = message
 
@@ -102,35 +108,32 @@ def rate_limit_percent(self) -> float:
 
         return self.rate_limit_remaining / self.rate_limit * 100
 
-    def agencies(self) -> list[dict[str, Any]]:
-        result = self._get("agencies")
-        return cast(list[dict[str, Any]], result)
+    def agencies(self) -> list[AgencyInfo]:
+        return cast(list[AgencyInfo], self._get("agencies"))
 
-    def routes(self, agency_id: str | None = None) -> list[dict[str, Any]]:
+    def routes(self, agency_id: str | None = None) -> list[RouteInfo]:
         if not agency_id:
             agency_id = self.agency_id
 
-        result = self._get(f"agencies/{agency_id}/routes")
-        return cast(list[dict[str, Any]], result)
+        return cast(list[RouteInfo], self._get(f"agencies/{agency_id}/routes"))
 
     def route_details(
         self, route_id: str, agency_id: str | None = None
-    ) -> dict[str, Any] | str:
+    ) -> RouteDetails:
         """Includes stops and directions."""
         agency_id = agency_id or self.agency_id
         if not agency_id:
             raise NextBusValidationError("Agency ID is required")
 
-        result = self._get(f"agencies/{agency_id}/routes/{route_id}")
-        return cast(dict[str, Any], result)
+        return cast(RouteDetails, self._get(f"agencies/{agency_id}/routes/{route_id}"))
 
     def predictions_for_stop(
         self,
         stop_id: str | int,
         route_id: str | None = None,
         direction_id: str | None = None,
         agency_id: str | None = None,
-    ) -> list[dict[str, Any]]:
+    ) -> list[StopPrediction]:
         """Returns predictions for a stop."""
         agency_id = agency_id or self.agency_id
         if not agency_id:
@@ -141,13 +144,17 @@ def predictions_for_stop(
                 raise NextBusValidationError("Direction ID provided without route ID")
 
         if route_id:
-            result = self._get(
-                f"agencies/{agency_id}/nstops/{route_id}:{stop_id}/predictions"
+            predictions = cast(
+                list[StopPrediction],
+                self._get(
+                    f"agencies/{agency_id}/nstops/{route_id}:{stop_id}/predictions"
+                ),
             )
         else:
-            result = self._get(f"agencies/{agency_id}/stops/{stop_id}/predictions")
-
-        predictions = cast(list[dict[str, Any]], result)
+            predictions = cast(
+                list[StopPrediction],
+                self._get(f"agencies/{agency_id}/stops/{stop_id}/predictions"),
+            )
 
         # If route not provided, return all predictions as the API returned them
         if not route_id:
@@ -174,9 +181,7 @@ def predictions_for_stop(
 
         return predictions
 
-    def _get(
-        self, endpoint: str, params: dict[str, Any] | None = None
-    ) -> dict[str, Any] | list[dict[str, Any]]:
+    def _get(self, endpoint: str, params: dict[str, Any] | None = None) -> Any:
         if params is None:
             params = {}
 

py_nextbus/models.py

@@ -0,0 +1,119 @@
+from typing import TypedDict
+
+
+class AgencyInfo(TypedDict):
+    id: str
+    name: str
+    shortName: str
+    region: str
+    website: str
+    logo: str
+    nxbs2RedirectUrl: str
+
+
+class PredictionRouteInfo(TypedDict):
+    id: str
+    title: str
+    description: str
+    color: str
+    textColor: str
+    hidden: bool
+
+
+class RouteInfo(PredictionRouteInfo):
+    rev: int
+    timestamp: str
+
+
+class RouteBoundingBox(TypedDict):
+    latMin: float
+    latMax: float
+    lonMin: float
+    lonMax: float
+
+
+class StopInfo(TypedDict):
+    id: str
+    lat: float
+    lon: float
+    name: str
+    code: str
+    hidden: bool
+    showDestinationSelector: bool
+    directions: list[str]
+
+
+class PredictionStopInfo(TypedDict):
+    id: str
+    lat: float
+    lon: float
+    name: str
+    code: str
+    hidden: bool
+    showDestinationSelector: bool
+    route: str
+
+
+class Point(TypedDict):
+    lat: float
+    lon: float
+
+
+class RoutePath(TypedDict):
+    id: str
+    points: list[Point]
+
+
+class DirectionInfo(TypedDict):
+    id: str
+    shortName: str
+    name: str
+    useForUi: bool
+    stops: list[str]
+
+
+class RouteDetails(TypedDict):
+    id: str
+    rev: int
+    title: str
+    description: str
+    color: str
+    textColor: str
+    hidden: bool
+    boundingBox: RouteBoundingBox
+    stops: list[StopInfo]
+    directions: list[DirectionInfo]
+    paths: list[RoutePath]
+    timestamp: str
+
+
+class PredictionDirection(TypedDict):
+    id: str
+    name: str
+    destinationName: str
+
+
+class PredictionValue(TypedDict):
+    timestamp: int
+    minutes: int
+    affectedByLayover: bool
+    isDeparture: bool
+    occupancyStatus: int
+    occupancyDescription: str
+    vehiclesInConsist: int
+    linkedVehicleIds: str
+    vehicleId: str
+    vehicleType: str | None
+    direction: PredictionDirection
+    tripId: str
+    delay: int
+    predUsingNavigationTm: bool
+    departure: bool
+
+
+class StopPrediction(TypedDict):
+    serverTimestamp: int
+    nxbs2RedirectUrl: str
+    route: PredictionRouteInfo
+    stop: PredictionStopInfo
+    values: list[PredictionValue]

tests/client_test.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import unittest.mock
+from unittest.mock import MagicMock
 
 from py_nextbus.client import NextBusClient
 from tests.helpers.mock_responses import MOCK_PREDICTIONS_RESPONSE_NO_ROUTE
@@ -16,7 +17,7 @@ def setUp(self):
         self.client = NextBusClient()
 
     @unittest.mock.patch("py_nextbus.client.NextBusClient._get")
-    def test_predictions_for_stop_no_route(self, mock_get):
+    def test_predictions_for_stop_no_route(self, mock_get: MagicMock):
         mock_get.return_value = MOCK_PREDICTIONS_RESPONSE_NO_ROUTE
 
         result = self.client.predictions_for_stop(
@@ -32,7 +33,7 @@ def test_predictions_for_stop_no_route(self, mock_get):
         )
 
     @unittest.mock.patch("py_nextbus.client.NextBusClient._get")
-    def test_predictions_for_stop_with_route(self, mock_get):
+    def test_predictions_for_stop_with_route(self, mock_get: MagicMock):
         mock_get.return_value = MOCK_PREDICTIONS_RESPONSE_WITH_ROUTE
 
         result = self.client.predictions_for_stop(

tests/helpers/mock_responses.py

@@ -1,9 +1,14 @@
+from py_nextbus.models import AgencyInfo
+from py_nextbus.models import RouteDetails
+from py_nextbus.models import RouteInfo
+from py_nextbus.models import StopPrediction
+
 TEST_AGENCY_ID = "sfmta-cis"
 TEST_ROUTE_ID = "F"
 TEST_STOP_ID = "5184"
 TEST_DIRECTION_ID = "F_0_var0"
 
-MOCK_AGENCY_LIST_RESPONSE = [
+MOCK_AGENCY_LIST_RESPONSE: list[AgencyInfo] = [
     {
         "id": "sfmta-cis",
         "name": "San Francisco Muni CIS",
@@ -15,7 +20,7 @@
     },
 ]
 
-MOCK_ROUTE_LIST_RESPONSE = [
+MOCK_ROUTE_LIST_RESPONSE: list[RouteInfo] = [
     {
         "id": "F",
         "rev": 1057,
@@ -28,7 +33,7 @@
     },
 ]
 
-MOCK_ROUTE_DETAILS_RESPONSE = {
+MOCK_ROUTE_DETAILS_RESPONSE: RouteDetails = {
     "id": "F",
     "rev": 1057,
     "title": "F Market & Wharves",
@@ -105,7 +110,7 @@
     "timestamp": "2024-06-23T03:06:58Z",
 }
 
-MOCK_PREDICTIONS_RESPONSE_NO_ROUTE = [
+MOCK_PREDICTIONS_RESPONSE_NO_ROUTE: list[StopPrediction] = [
     {
         "serverTimestamp": 1724038210798,
         "nxbs2RedirectUrl": "",
@@ -241,7 +246,7 @@
     },
 ]
 
-MOCK_PREDICTIONS_RESPONSE_WITH_ROUTE = [
+MOCK_PREDICTIONS_RESPONSE_WITH_ROUTE: list[StopPrediction] = [
     {
         "serverTimestamp": 1720034290432,
         "nxbs2RedirectUrl": "",