From 833a4785d4408cf16e52a26a80a8fb1e96854e45 Mon Sep 17 00:00:00 2001 From: Lex Alexander Date: Fri, 4 Sep 2026 13:14:18 -0400 Subject: [PATCH 1/5] Harden MARTA API requests and credential handling --- marta/api.py | 109 +++++++++++++++++------ marta/exceptions.py | 4 +- pytest.ini | 3 + requirements.txt | 4 +- setup.py | 6 +- tests/test_api.py | 208 ++++++++++++++++++++++++++++++++------------ 6 files changed, 241 insertions(+), 93 deletions(-) create mode 100644 pytest.ini diff --git a/marta/api.py b/marta/api.py index fda7e35..df8739e 100644 --- a/marta/api.py +++ b/marta/api.py @@ -1,21 +1,29 @@ +import json +from functools import wraps +from os import getenv from typing import List, Union -import requests import requests_cache -import json -from os import getenv -from functools import wraps from .exceptions import APIKeyError, InvalidDirectionError from .vehicles import Bus, Train _CACHE_EXPIRE = int(getenv('MARTA_CACHE_EXPIRE', 30)) -_BASE_URL = 'http://developer.itsmarta.com' +_BASE_URL = 'https://developer.itsmarta.com' _TRAIN_PATH = '/RealtimeTrain/RestServiceNextTrain/GetRealtimeArrivals' _BUS_PATH = '/BRDRestService/RestBusRealTimeService/GetAllBus' _BUS_ROUTE_PATH = '/BRDRestService/RestBusRealTimeService/GetBusByRoute/' - -requests_cache.install_cache('marta_api_cache', backend='sqlite', expire_after=_CACHE_EXPIRE) +_REQUEST_TIMEOUT = (3.05, 10) +_MAX_RESPONSE_BYTES = 5 * 1024 * 1024 +_JSON_CONTENT_TYPES = {'application/json', 'text/json'} + +# Keep caching private to this library and in memory. Importing marta must not +# patch requests globally or persist request URLs containing credentials. +CACHE = requests_cache.CachedSession( + backend='memory', + expire_after=_CACHE_EXPIRE, + ignored_parameters=['apikey'], +) def require_api_key(func): """ @@ -31,41 +39,84 @@ def with_key(self, *args, **kwargs): return with_key +def get_bus_direction(user_direction) -> str: + direction = user_direction.lower() + if direction.startswith('n'): + return 'Northbound' + elif direction.startswith('s'): + return 'Southbound' + elif direction.startswith('e'): + return 'Eastbound' + elif direction.startswith('w'): + return 'Westbound' + else: + return None + +def get_train_direction(user_direction) -> str: + train_direction = user_direction.lower() + if train_direction.startswith('n'): + return 'N' + elif train_direction.startswith('s'): + return 'S' + elif train_direction.startswith('e'): + return 'E' + elif train_direction.startswith('w'): + return 'W' + else: + return None + + def _convert_direction(user_direction: str, vehicle_type: str = 'bus') -> Union[str, None]: if not user_direction: return None if vehicle_type == 'bus': - if user_direction.lower().startswith('n'): - return 'Northbound' - elif user_direction.lower().startswith('s'): - return 'Southbound' - elif user_direction.lower().startswith('e'): - return 'Eastbound' - elif user_direction.lower().startswith('w'): - return 'Westbound' + bus_direction = get_bus_direction(user_direction) + if bus_direction is not None: + return bus_direction else: raise InvalidDirectionError(direction_provided=user_direction) elif vehicle_type == 'train': - if user_direction.lower().startswith('n'): - return 'N' - elif user_direction.lower().startswith('s'): - return 'S' - elif user_direction.lower().startswith('e'): - return 'E' - elif user_direction.lower().startswith('w'): - return 'W' + train_direction = get_train_direction(user_direction) + if train_direction is not None: + return train_direction else: raise InvalidDirectionError(direction_provided=user_direction) else: return user_direction -def _get_data(endpoint: str, api_key: str) -> dict: - url = f'{_BASE_URL}{endpoint}?apikey={api_key}' - response = requests.get(url) - if response.status_code == 401 or response.status_code == 403: - raise APIKeyError(f'Your API key seems to be invalid. Try visiting {url}.') - return json.loads(response.text) +def _get_data(endpoint: str, api_key: str) -> List[dict]: + url = f'{_BASE_URL}{endpoint}' + response = CACHE.get( + url, + params={'apikey': api_key}, + timeout=_REQUEST_TIMEOUT, + stream=True, + ) + if response.status_code in (401, 403): + raise APIKeyError('Your MARTA API key was rejected.') + response.raise_for_status() + + content_type = response.headers.get('Content-Type', '').split(';', 1)[0].lower() + if content_type and content_type not in _JSON_CONTENT_TYPES: + raise ValueError(f'Unexpected MARTA API content type: {content_type}') + + content_length = response.headers.get('Content-Length') + if content_length and int(content_length) > _MAX_RESPONSE_BYTES: + raise ValueError('MARTA API response exceeds the maximum allowed size') + + chunks = [] + received = 0 + for chunk in response.iter_content(chunk_size=64 * 1024): + received += len(chunk) + if received > _MAX_RESPONSE_BYTES: + raise ValueError('MARTA API response exceeds the maximum allowed size') + chunks.append(chunk) + + data = json.loads(b''.join(chunks).decode(response.encoding or 'utf-8')) + if not isinstance(data, list) or any(not isinstance(item, dict) for item in data): + raise ValueError('MARTA API response must be a list of objects') + return data def _filter_response(response: dict, filters: dict) -> List[dict]: valid_items = [] diff --git a/marta/exceptions.py b/marta/exceptions.py index dab3668..65a08de 100644 --- a/marta/exceptions.py +++ b/marta/exceptions.py @@ -4,7 +4,7 @@ def __init__(self, message: str = None): if not message: message = 'API Key is missing. Please set MARTA_API_KEY or use api_key kwarg.' - super(Exception, self).__init__(message) + super().__init__(message) class InvalidDirectionError(Exception): """Exception thrown for an invalid bus/train direction""" @@ -12,4 +12,4 @@ def __init__(self, direction_provided: str, message: str = None): if not message: message = f'{direction_provided} is an invalid direction.' - super(Exception, self).__init__(message) \ No newline at end of file + super().__init__(message) diff --git a/pytest.ini b/pytest.ini new file mode 100644 index 0000000..c7b23ec --- /dev/null +++ b/pytest.ini @@ -0,0 +1,3 @@ +[pytest] +pythonpath = . +testpaths = tests diff --git a/requirements.txt b/requirements.txt index 6493c2e..9a6e556 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,2 +1,2 @@ -requests -requests-cache +requests>=2.32.4,<3 +requests-cache>=1.2,<2 diff --git a/setup.py b/setup.py index fac285d..97db2bf 100644 --- a/setup.py +++ b/setup.py @@ -3,11 +3,11 @@ setup( name='marta', description='Python library for accessing MARTA real-time API', - url='http://www.itsmarta.com/app-developer-resources.aspx', + url='https://www.itsmarta.com/app-developer-resources.aspx', packages=['marta'], install_requires=[ 'pytest-runner', - 'requests', - 'requests-cache' + 'requests>=2.32.4,<3', + 'requests-cache>=1.2,<2' ], ) diff --git a/tests/test_api.py b/tests/test_api.py index f35e453..bd27ff7 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -1,82 +1,176 @@ +import json + import pytest -import requests_mock +import requests -import marta -from marta.api import get_buses, get_trains +from marta.api import MARTA, _get_data from marta.exceptions import APIKeyError from marta.vehicles import Bus, Train -def test_get_trains(train_response): - with requests_mock.mock() as m: - m.get(requests_mock.ANY, text=train_response) - trains = get_trains() - for t in trains: - assert isinstance(t, Train) +class FakeResponse: + def __init__(self, payload=None, status_code=200, headers=None, chunks=None): + self.payload = payload + self.status_code = status_code + self.headers = ( + {'Content-Type': 'application/json'} if headers is None else headers + ) + self.encoding = 'utf-8' + self._chunks = chunks + + def raise_for_status(self): + if self.status_code >= 400: + raise requests.HTTPError() + + def iter_content(self, chunk_size): + if self._chunks is not None: + yield from self._chunks + return + yield json.dumps(self.payload).encode() + + +@pytest.fixture +def mock_response(monkeypatch): + calls = [] + + def install(response): + def get(url, **kwargs): + calls.append((url, kwargs)) + return response + + monkeypatch.setattr('marta.api.CACHE.get', get) + return calls + + return install + + +def test_get_trains(mock_response, train_response): + calls = mock_response(FakeResponse(json.loads(train_response))) + + trains = MARTA(api_key='secret').get_trains() + + assert trains + assert all(isinstance(train, Train) for train in trains) + url, kwargs = calls[0] + assert url == ( + 'https://developer.itsmarta.com/' + 'RealtimeTrain/RestServiceNextTrain/GetRealtimeArrivals' + ) + assert 'secret' not in url + assert kwargs == { + 'params': {'apikey': 'secret'}, + 'timeout': (3.05, 10), + 'stream': True, + } + + +@pytest.mark.parametrize( + ('filters', 'attribute', 'expected'), + [ + ({'line': 'BLUE'}, 'line', 'BLUE'), + ({'station': 'INDIAN CREEK STATION'}, 'station', 'INDIAN CREEK STATION'), + ], +) +def test_get_trains_filters_locally( + mock_response, train_response, filters, attribute, expected +): + mock_response(FakeResponse(json.loads(train_response))) + + trains = MARTA(api_key='secret').get_trains(**filters) + + assert trains + assert all(getattr(train, attribute) == expected for train in trains) + + +def test_get_trains_combines_filters(mock_response, train_response): + mock_response(FakeResponse(json.loads(train_response))) + + trains = MARTA(api_key='secret').get_trains( + line='BLUE', station='INDIAN CREEK STATION' + ) + + assert len(trains) == 1 + assert trains[0].line == 'BLUE' + assert trains[0].station == 'INDIAN CREEK STATION' + + +def test_get_buses(mock_response, bus_all_response): + calls = mock_response(FakeResponse(json.loads(bus_all_response))) + + buses = MARTA(api_key='secret').get_buses() + + assert buses + assert all(isinstance(bus, Bus) for bus in buses) + assert calls[0][0] == ( + 'https://developer.itsmarta.com/' + 'BRDRestService/RestBusRealTimeService/GetAllBus' + ) + + +def test_get_buses_by_route(mock_response, bus_route_response): + mock_response(FakeResponse(json.loads(bus_route_response))) + + buses = MARTA(api_key='secret').get_buses(route=1) + + assert buses + assert all(bus.route == 1 for bus in buses) + + +def test_missing_api_key_is_deterministic(monkeypatch): + monkeypatch.delenv('MARTA_API_KEY', raising=False) + + with pytest.raises(APIKeyError, match='API Key is missing'): + MARTA().get_buses() -def test_get_trains_by_line(train_response): - line = "BLUE" - with requests_mock.mock() as m: - m.get(requests_mock.ANY, text=train_response) - trains = get_trains(line=line) +@pytest.mark.parametrize('status_code', [401, 403]) +def test_rejected_key_is_not_disclosed(mock_response, status_code): + mock_response(FakeResponse(status_code=status_code)) - assert len(trains) > 0 - for t in trains: - assert isinstance(t, Train) - assert t.line == line + with pytest.raises(APIKeyError) as error: + _get_data('/endpoint', 'super-secret') + assert 'super-secret' not in str(error.value) + assert str(error.value) == 'Your MARTA API key was rejected.' -def test_get_trains_by_station(train_response): - station = "INDIAN CREEK STATION" - with requests_mock.mock() as m: - m.get(requests_mock.ANY, text=train_response) - trains = get_trains(station=station) - assert len(trains) > 0 - for t in trains: - assert isinstance(t, Train) - assert t.station == station +def test_raises_for_other_http_errors(mock_response): + mock_response(FakeResponse(status_code=500)) + with pytest.raises(requests.HTTPError): + _get_data('/endpoint', 'secret') -def test_get_trains_by_line_and_station(train_response): - station = "INDIAN CREEK STATION" - line = "BLUE" - with requests_mock.mock() as m: - m.get(requests_mock.ANY, text=train_response) - trains = get_trains(line=line, station=station) - assert len(trains) > 0 - for t in trains: - assert isinstance(t, Train) - assert t.station == station - assert t.line == line +def test_rejects_non_json_response(mock_response): + mock_response(FakeResponse(headers={'Content-Type': 'text/html'})) + with pytest.raises(ValueError, match='content type'): + _get_data('/endpoint', 'secret') -def test_get_buses(bus_all_response): - with requests_mock.mock() as m: - m.get(requests_mock.ANY, text=bus_all_response) - buses = get_buses() - assert len(buses) > 0 +def test_rejects_oversized_content_length(mock_response): + mock_response( + FakeResponse( + headers={ + 'Content-Type': 'application/json', + 'Content-Length': str(5 * 1024 * 1024 + 1), + } + ) + ) - for b in buses: - assert isinstance(b, Bus) + with pytest.raises(ValueError, match='maximum allowed size'): + _get_data('/endpoint', 'secret') -def test_get_buses_by_route(bus_route_response): - with requests_mock.mock() as m: - m.get(requests_mock.ANY, text=bus_route_response) - buses = get_buses(route=1) +def test_rejects_oversized_stream(mock_response): + mock_response(FakeResponse(chunks=[b'x' * (5 * 1024 * 1024 + 1)])) - assert len(buses) > 0 + with pytest.raises(ValueError, match='maximum allowed size'): + _get_data('/endpoint', 'secret') - for b in buses: - assert isinstance(b, Bus) - assert b.route == 1 +def test_rejects_unexpected_json_shape(mock_response): + mock_response(FakeResponse({'not': 'a list'})) -def test_missing_api_key(): - marta.api._API_KEY = None - with pytest.raises(APIKeyError): - buses = get_buses() + with pytest.raises(ValueError, match='list of objects'): + _get_data('/endpoint', 'secret') From e9ef12aa18ce55fa6bb1e1e7f077d582f0518fec Mon Sep 17 00:00:00 2001 From: Lex Alexander Date: Fri, 4 Sep 2026 15:04:45 -0400 Subject: [PATCH 2/5] Add GTFS for bus routes and upgrade train urls --- README.md | 60 +++++++------ marta/api.py | 140 +++++++++++++++++++++--------- marta/vehicles.py | 72 +++++++++++++++- requirements.txt | 3 +- setup.py | 3 +- tests/test_api.py | 189 ++++++++++++++++++++++++++++++++++++----- tests/test_vehicles.py | 32 +++++++ 7 files changed, 414 insertions(+), 85 deletions(-) diff --git a/README.md b/README.md index 36a78d6..f4d9fed 100644 --- a/README.md +++ b/README.md @@ -47,46 +47,55 @@ Mac/Linux: export MARTA_CACHE_EXPIRE=15 ``` -There are two primary API wrapper functions, `get_buses()` and `get_trains()`. Each method takes keyword arguments to filter results. +Create a client and use `get_buses()` or `get_trains()` for the existing object +interface. `get_buses()` now reads MARTA's GTFS-Realtime vehicle feed and +adapts each vehicle position into a `Bus` object. + +> **Deprecation warning:** MARTA removed its legacy bus JSON endpoints. The +> `get_buses()` arguments `time_point`, `direction`, and `api_key`, plus the +> `Bus` attributes `adherence`, `block_id`, `block_abbr`, `direction`, and +> `timepoint`, are retained only for compatibility. GTFS-Realtime does not +> provide those values, so the attributes are `None` and the unsupported +> filters raise `ValueError`. New integrations should use +> `get_bus_vehicle_positions()` for the source protobuf feed. ``` -from marta.api import get_buses, get_trains +from marta import MARTA + +marta = MARTA() # Get all buses -buses = get_buses() +buses = marta.get_buses() # Get buses by route -buses = get_buses(route=1) +buses = marta.get_buses(route=1) # Get buses by route and stop_id -buses = get_buses(route=1, stop_id=900800) +buses = marta.get_buses(route=1, stop_id=900800) # Get buses by route and vehicle_id -buses = get_buses(route=1, bus_id=1405) - -# Get buses by route and timepoint -buses = get_buses(route=1, timepoint="West End Station") +buses = marta.get_buses(route=1, vehicle_id=1405) # Get buses by route, stop_id and vehicle_id -buses = get_buses(route=1, stop_id=900800, bus_id=1405) +buses = marta.get_buses(route=1, stop_id=900800, vehicle_id=1405) # Get all trains -trains = get_trains() +trains = marta.get_trains() # Get trains by line -trains = get_trains(line='red') +trains = marta.get_trains(line='red') # Get trains by station -trains = get_trains(station='Midtown Station') +trains = marta.get_trains(station='Midtown Station') # Get trains by destination -trains = get_trains(destination='Doraville') +trains = marta.get_trains(destination='Doraville') # Get trains by line, station, and destination -trains = get_trains(line='blue', station='Five Points Station', destination='Indian Creek') +trains = marta.get_trains(line='blue', station='Five Points Station', destination='Indian Creek') # Get trains by line, station, and direction -trains = get_trains(line='gold', station='Five Points Station', direction='north') +trains = marta.get_trains(line='gold', station='Five Points Station', direction='north') ``` ## Results @@ -98,17 +107,18 @@ The `get_buses()` and `get_trains()` functions return lists of `Bus` and `Train` ### Bus Objects ``` -{'adherence': '-3', - 'block_abbr': '3-3', - 'block_id': '346', - 'direction': 'Westbound', - 'last_updated': datetime.datetime(2017, 1, 26, 8, 10, 24), - 'latitude': '33.7545535', - 'longitude': '-84.4686002', +{'adherence': None, + 'block_abbr': None, + 'block_id': None, + 'direction': None, + 'direction_id': 0, + 'last_updated': datetime.datetime(2026, 9, 4, 17, 0, tzinfo=datetime.timezone.utc), + 'latitude': 33.7545535, + 'longitude': -84.4686002, 'route': 3, 'stop_id': '903320', - 'timepoint': 'Hamilton E. Holmes Station', - 'trip_id': '5408210', + 'timepoint': None, + 'trip_id': 'trip-1', 'vehicle': '2417'} ``` diff --git a/marta/api.py b/marta/api.py index df8739e..0128a4c 100644 --- a/marta/api.py +++ b/marta/api.py @@ -2,27 +2,44 @@ from functools import wraps from os import getenv from typing import List, Union +from urllib.parse import urlparse import requests_cache +from google.transit import gtfs_realtime_pb2 from .exceptions import APIKeyError, InvalidDirectionError from .vehicles import Bus, Train _CACHE_EXPIRE = int(getenv('MARTA_CACHE_EXPIRE', 30)) _BASE_URL = 'https://developer.itsmarta.com' -_TRAIN_PATH = '/RealtimeTrain/RestServiceNextTrain/GetRealtimeArrivals' -_BUS_PATH = '/BRDRestService/RestBusRealTimeService/GetAllBus' -_BUS_ROUTE_PATH = '/BRDRestService/RestBusRealTimeService/GetBusByRoute/' +_TRAIN_URL = ( + 'https://developerservices.itsmarta.com:18096/itsmarta/' + 'railrealtimearrivals/developerservices/traindata' +) +_BUS_VEHICLE_POSITIONS_URL = ( + 'https://gtfs-rt.itsmarta.com/TMGTFSRealTimeWebService/' + 'vehicle/vehiclepositions.pb' +) +_BUS_TRIP_UPDATES_URL = ( + 'https://gtfs-rt.itsmarta.com/TMGTFSRealTimeWebService/' + 'tripupdate/tripupdates.pb' +) _REQUEST_TIMEOUT = (3.05, 10) _MAX_RESPONSE_BYTES = 5 * 1024 * 1024 _JSON_CONTENT_TYPES = {'application/json', 'text/json'} +_PROTOBUF_CONTENT_TYPES = { + 'application/octet-stream', + 'application/protocol-buffer', + 'application/protobuf', + 'application/x-protobuf', +} # Keep caching private to this library and in memory. Importing marta must not # patch requests globally or persist request URLs containing credentials. CACHE = requests_cache.CachedSession( backend='memory', expire_after=_CACHE_EXPIRE, - ignored_parameters=['apikey'], + ignored_parameters=['apikey', 'apiKey'], ) def require_api_key(func): @@ -85,20 +102,11 @@ def _convert_direction(user_direction: str, vehicle_type: str = 'bus') -> Union[ return user_direction -def _get_data(endpoint: str, api_key: str) -> List[dict]: - url = f'{_BASE_URL}{endpoint}' - response = CACHE.get( - url, - params={'apikey': api_key}, - timeout=_REQUEST_TIMEOUT, - stream=True, - ) - if response.status_code in (401, 403): - raise APIKeyError('Your MARTA API key was rejected.') +def _read_response(response, allowed_content_types) -> bytes: response.raise_for_status() content_type = response.headers.get('Content-Type', '').split(';', 1)[0].lower() - if content_type and content_type not in _JSON_CONTENT_TYPES: + if content_type and content_type not in allowed_content_types: raise ValueError(f'Unexpected MARTA API content type: {content_type}') content_length = response.headers.get('Content-Length') @@ -112,12 +120,55 @@ def _get_data(endpoint: str, api_key: str) -> List[dict]: if received > _MAX_RESPONSE_BYTES: raise ValueError('MARTA API response exceeds the maximum allowed size') chunks.append(chunk) + return b''.join(chunks) + - data = json.loads(b''.join(chunks).decode(response.encoding or 'utf-8')) +def _get_json_data( + url: str, + api_key: str, + api_key_parameter: str, +) -> List[dict]: + response = CACHE.get( + url, + params={api_key_parameter: api_key}, + timeout=_REQUEST_TIMEOUT, + stream=True, + ) + if response.status_code in (401, 403): + raise APIKeyError('Your MARTA API key was rejected.') + content = _read_response(response, _JSON_CONTENT_TYPES) + data = json.loads(content.decode(response.encoding or 'utf-8')) if not isinstance(data, list) or any(not isinstance(item, dict) for item in data): raise ValueError('MARTA API response must be a list of objects') return data + +def _get_data(endpoint: str, api_key: str) -> List[dict]: + """Fetch a legacy MARTA JSON endpoint.""" + return _get_json_data( + url=f'{_BASE_URL}{endpoint}', + api_key=api_key, + api_key_parameter='apikey', + ) + + +def _get_gtfs_realtime_feed(url: str) -> gtfs_realtime_pb2.FeedMessage: + response = CACHE.get( + url, + timeout=_REQUEST_TIMEOUT, + stream=True, + allow_redirects=False, + ) + if response.is_redirect or response.is_permanent_redirect: + redirect_url = response.headers.get('Location', '') + if urlparse(redirect_url).scheme != 'https': + raise ValueError('MARTA API attempted an insecure redirect') + raise ValueError('MARTA API returned an unexpected redirect') + content = _read_response(response, _PROTOBUF_CONTENT_TYPES) + feed = gtfs_realtime_pb2.FeedMessage() + feed.ParseFromString(content) + return feed + def _filter_response(response: dict, filters: dict) -> List[dict]: valid_items = [] for item in response: @@ -161,7 +212,11 @@ def get_trains(self, :return: list of Train objects :rtype: List[Train] """ - data = _get_data(endpoint=_TRAIN_PATH, api_key=api_key) + data = _get_json_data( + url=_TRAIN_URL, + api_key=api_key, + api_key_parameter='apiKey', + ) filters = { 'LINE': line, 'DIRECTION': _convert_direction(user_direction=direction, vehicle_type='train'), @@ -171,44 +226,53 @@ def get_trains(self, matching_data = _filter_response(response=data, filters=filters) return [Train(t) for t in matching_data] - @require_api_key def get_buses(self, route: int = None, stop_id: int = None, vehicle_id: int = None, time_point: str = None, - direction: str = None, - api_key: str = None) -> List[Bus]: + direction: str = None + ) -> List[Bus]: """ - Query API for bus information + Query the GTFS-Realtime vehicle feed for bus information. :param route: route number :type route: int, optional :param stop_id: Bus stop ID :type stop_id: int, optional :param vehicle_id: Bus ID :type vehicle_id: int, optional - :param time_point: + :param time_point: Deprecated; unavailable in GTFS-Realtime :type time_point: str, optional - :param direction: Bus direction (Northbound, Southbound, Westbound or Eastbound) + :param direction: Deprecated; cardinal direction is unavailable in GTFS-Realtime :type direction: str, optional - :param api_key: API key to override environment variable + :param api_key: Deprecated; the GTFS-Realtime feed does not require a key :type api_key: str, optional :return: list of Bus objects """ - if route: - endpoint = f'{_BUS_ROUTE_PATH}/{route}' - else: - endpoint = f'{_BUS_PATH}' + if time_point is not None or direction is not None: + raise ValueError( + 'time_point and direction filters are unavailable in ' + 'MARTA GTFS-Realtime data' + ) + feed = self.get_bus_vehicle_positions() + buses = [ + Bus.from_gtfs(entity.vehicle) + for entity in feed.entity + if entity.HasField('vehicle') + ] + return [ + bus for bus in buses + if (route is None or str(bus.route) == str(route)) + and (stop_id is None or str(bus.stop_id) == str(stop_id)) + and (vehicle_id is None or str(bus.vehicle) == str(vehicle_id)) + ] - data = _get_data(endpoint=endpoint, api_key=api_key) - filters = { - 'STOPID': stop_id, - 'VEHICLE': vehicle_id, - 'TIMEPOINT': time_point, - 'ROUTE': route, - 'DIRECTION': _convert_direction(user_direction=direction, vehicle_type='bus') - } - matching_data = _filter_response(response=data, filters=filters) - return [Bus(b) for b in matching_data] + def get_bus_vehicle_positions(self) -> gtfs_realtime_pb2.FeedMessage: + """Return MARTA's GTFS-Realtime bus vehicle-position feed.""" + return _get_gtfs_realtime_feed(_BUS_VEHICLE_POSITIONS_URL) + + def get_bus_trip_updates(self) -> gtfs_realtime_pb2.FeedMessage: + """Return MARTA's GTFS-Realtime bus trip-update feed.""" + return _get_gtfs_realtime_feed(_BUS_TRIP_UPDATES_URL) diff --git a/marta/vehicles.py b/marta/vehicles.py index 00dcbbb..02cb6da 100644 --- a/marta/vehicles.py +++ b/marta/vehicles.py @@ -1,4 +1,4 @@ -from datetime import datetime +from datetime import datetime, timezone class Vehicle(): @@ -23,6 +23,76 @@ def __init__(self, record): self.trip_id = record.get('TRIPID') self.vehicle = record.get('VEHICLE') + @classmethod + def from_gtfs(cls, vehicle_position): + """Build a compatibility Bus from a GTFS-Realtime VehiclePosition.""" + bus = cls.__new__(cls) + trip = vehicle_position.trip + vehicle = vehicle_position.vehicle + position = vehicle_position.position + + bus.adherence = None + bus.block_id = None + bus.block_abbr = None + bus.direction = None + bus.direction_id = ( + trip.direction_id if trip.HasField('direction_id') else None + ) + bus.latitude = position.latitude + bus.longitude = position.longitude + bus.bearing = position.bearing if position.HasField('bearing') else None + bus.speed = position.speed if position.HasField('speed') else None + bus.last_updated = ( + datetime.fromtimestamp(vehicle_position.timestamp, timezone.utc) + if vehicle_position.HasField('timestamp') + else None + ) + bus.route = _coerce_route(trip.route_id) + bus.stop_id = vehicle_position.stop_id or None + bus.timepoint = None + bus.trip_id = trip.trip_id or None + bus.vehicle = vehicle.id or None + bus.current_status = ( + vehicle_position.current_status + if vehicle_position.HasField('current_status') + else None + ) + bus.raw_data = { + 'trip': { + 'trip_id': bus.trip_id, + 'route_id': trip.route_id or None, + 'direction_id': bus.direction_id, + }, + 'vehicle': { + 'id': bus.vehicle, + 'label': vehicle.label or None, + 'license_plate': vehicle.license_plate or None, + }, + 'position': { + 'latitude': bus.latitude, + 'longitude': bus.longitude, + 'bearing': bus.bearing, + 'speed': bus.speed, + }, + 'stop_id': bus.stop_id, + 'timestamp': ( + vehicle_position.timestamp + if vehicle_position.HasField('timestamp') + else None + ), + 'current_status': bus.current_status, + } + return bus + + +def _coerce_route(route_id): + if not route_id: + return None + try: + return int(route_id) + except ValueError: + return route_id + class Train(Vehicle): def __init__(self, record): diff --git a/requirements.txt b/requirements.txt index 9a6e556..e97c23e 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,2 +1,3 @@ +gtfs-realtime-bindings>=2.2,<3 requests>=2.32.4,<3 -requests-cache>=1.2,<2 +requests-cache>=1.2,<2 \ No newline at end of file diff --git a/setup.py b/setup.py index 97db2bf..36c532a 100644 --- a/setup.py +++ b/setup.py @@ -5,8 +5,9 @@ description='Python library for accessing MARTA real-time API', url='https://www.itsmarta.com/app-developer-resources.aspx', packages=['marta'], + python_requires='>=3.8', install_requires=[ - 'pytest-runner', + 'gtfs-realtime-bindings>=2.2,<3', 'requests>=2.32.4,<3', 'requests-cache>=1.2,<2' ], diff --git a/tests/test_api.py b/tests/test_api.py index bd27ff7..d826f52 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -2,7 +2,8 @@ import pytest import requests - +from google.protobuf.message import DecodeError +from google.transit import gtfs_realtime_pb2 from marta.api import MARTA, _get_data from marta.exceptions import APIKeyError from marta.vehicles import Bus, Train @@ -18,6 +19,14 @@ def __init__(self, payload=None, status_code=200, headers=None, chunks=None): self.encoding = 'utf-8' self._chunks = chunks + @property + def is_redirect(self): + return self.status_code in (301, 302, 303, 307, 308) + + @property + def is_permanent_redirect(self): + return self.status_code in (301, 308) + def raise_for_status(self): if self.status_code >= 400: raise requests.HTTPError() @@ -44,6 +53,22 @@ def get(url, **kwargs): return install +def gtfs_vehicle_feed(*vehicles): + feed = gtfs_realtime_pb2.FeedMessage() + feed.header.gtfs_realtime_version = '2.0' + for index, vehicle in enumerate(vehicles): + entity = feed.entity.add() + entity.id = f'vehicle-{index}' + entity.vehicle.trip.route_id = str(vehicle['route']) + entity.vehicle.trip.trip_id = vehicle.get('trip_id', f'trip-{index}') + entity.vehicle.vehicle.id = str(vehicle['vehicle_id']) + entity.vehicle.position.latitude = vehicle.get('latitude', 33.75) + entity.vehicle.position.longitude = vehicle.get('longitude', -84.39) + entity.vehicle.stop_id = str(vehicle['stop_id']) + entity.vehicle.timestamp = vehicle.get('timestamp', 1_725_000_000) + return feed + + def test_get_trains(mock_response, train_response): calls = mock_response(FakeResponse(json.loads(train_response))) @@ -53,15 +78,13 @@ def test_get_trains(mock_response, train_response): assert all(isinstance(train, Train) for train in trains) url, kwargs = calls[0] assert url == ( - 'https://developer.itsmarta.com/' - 'RealtimeTrain/RestServiceNextTrain/GetRealtimeArrivals' + 'https://developerservices.itsmarta.com:18096/itsmarta/' + 'railrealtimearrivals/developerservices/traindata' ) assert 'secret' not in url - assert kwargs == { - 'params': {'apikey': 'secret'}, - 'timeout': (3.05, 10), - 'stream': True, - } + assert kwargs['params'] == {'apiKey': 'secret'} + assert kwargs['timeout'] == (3.05, 10) + assert kwargs['stream'] == True @pytest.mark.parametrize( @@ -94,33 +117,161 @@ def test_get_trains_combines_filters(mock_response, train_response): assert trains[0].station == 'INDIAN CREEK STATION' -def test_get_buses(mock_response, bus_all_response): - calls = mock_response(FakeResponse(json.loads(bus_all_response))) +def test_get_buses_uses_gtfs_vehicle_feed(mock_response): + source = gtfs_vehicle_feed( + {'route': 1, 'vehicle_id': 1469, 'stop_id': 907473}, + {'route': 110, 'vehicle_id': 1463, 'stop_id': 901789}, + ) + calls = mock_response( + FakeResponse( + headers={'Content-Type': 'application/protocol-buffer'}, + chunks=[source.SerializeToString()], + ) + ) - buses = MARTA(api_key='secret').get_buses() + buses = MARTA().get_buses() - assert buses + assert len(buses) == 2 assert all(isinstance(bus, Bus) for bus in buses) assert calls[0][0] == ( - 'https://developer.itsmarta.com/' - 'BRDRestService/RestBusRealTimeService/GetAllBus' + 'https://gtfs-rt.itsmarta.com/TMGTFSRealTimeWebService/' + 'vehicle/vehiclepositions.pb' ) + assert 'params' not in calls[0][1] -def test_get_buses_by_route(mock_response, bus_route_response): - mock_response(FakeResponse(json.loads(bus_route_response))) +def test_get_buses_filters_gtfs_feed_locally(mock_response): + source = gtfs_vehicle_feed( + {'route': 1, 'vehicle_id': 1469, 'stop_id': 907473}, + {'route': 110, 'vehicle_id': 1463, 'stop_id': 901789}, + ) + mock_response( + FakeResponse( + headers={'Content-Type': 'application/x-protobuf'}, + chunks=[source.SerializeToString()], + ) + ) - buses = MARTA(api_key='secret').get_buses(route=1) + buses = MARTA().get_buses(route=1, stop_id=907473, vehicle_id=1469) - assert buses + assert len(buses) == 1 assert all(bus.route == 1 for bus in buses) + assert buses[0].stop_id == '907473' + assert buses[0].vehicle == '1469' + + +def test_get_bus_vehicle_positions_returns_gtfs_feed(mock_response): + source = gtfs_realtime_pb2.FeedMessage() + source.header.gtfs_realtime_version = '2.0' + entity = source.entity.add() + entity.id = 'vehicle-1' + entity.vehicle.vehicle.id = '1469' + calls = mock_response( + FakeResponse( + headers={'Content-Type': 'application/x-protobuf'}, + chunks=[source.SerializeToString()], + ) + ) + + feed = MARTA().get_bus_vehicle_positions() + + assert isinstance(feed, gtfs_realtime_pb2.FeedMessage) + assert feed.entity[0].vehicle.vehicle.id == '1469' + assert calls[0] == ( + ( + 'https://gtfs-rt.itsmarta.com/TMGTFSRealTimeWebService/' + 'vehicle/vehiclepositions.pb' + ), + {'timeout': (3.05, 10), 'stream': True, 'allow_redirects': False}, + ) + + +def test_get_bus_trip_updates_returns_gtfs_feed(mock_response): + source = gtfs_realtime_pb2.FeedMessage() + source.header.gtfs_realtime_version = '2.0' + entity = source.entity.add() + entity.id = 'trip-1' + entity.trip_update.trip.trip_id = '5391405' + calls = mock_response( + FakeResponse( + headers={'Content-Type': 'application/octet-stream'}, + chunks=[source.SerializeToString()], + ) + ) + + feed = MARTA().get_bus_trip_updates() + + assert isinstance(feed, gtfs_realtime_pb2.FeedMessage) + assert feed.entity[0].trip_update.trip.trip_id == '5391405' + assert calls[0] == ( + ( + 'https://gtfs-rt.itsmarta.com/TMGTFSRealTimeWebService/' + 'tripupdate/tripupdates.pb' + ), + {'timeout': (3.05, 10), 'stream': True, 'allow_redirects': False}, + ) + + +def test_get_buses_keeps_bus_object_contract(mock_response): + source = gtfs_vehicle_feed( + {'route': 1, 'vehicle_id': 1469, 'stop_id': 907473}, + ) + mock_response( + FakeResponse( + headers={'Content-Type': 'application/x-protobuf'}, + chunks=[source.SerializeToString()], + ) + ) + + buses = MARTA().get_buses() + + assert isinstance(buses, list) + assert all(isinstance(bus, Bus) for bus in buses) + assert buses[0].adherence is None + assert buses[0].timepoint is None + + +@pytest.mark.parametrize('filters', [{'time_point': 'Five Points'}, {'direction': 'N'}]) +def test_get_buses_rejects_unsupported_legacy_filters(filters): + with pytest.raises(ValueError, match='unavailable'): + MARTA().get_buses(**filters) + + +def test_rejects_invalid_protobuf(mock_response): + mock_response( + FakeResponse( + headers={'Content-Type': 'application/x-protobuf'}, + chunks=[b'not-a-protobuf-feed'], + ) + ) + + with pytest.raises(DecodeError): + MARTA().get_bus_vehicle_positions() + + +def test_rejects_insecure_gtfs_redirect(mock_response): + mock_response( + FakeResponse( + status_code=301, + headers={ + 'Content-Type': 'text/html', + 'Location': ( + 'http://gtfs-rt.itsmarta.com/' + 'TMGTFSRealTimeWebService/vehicle/' + ), + }, + ) + ) + + with pytest.raises(ValueError, match='insecure redirect'): + MARTA().get_bus_vehicle_positions() def test_missing_api_key_is_deterministic(monkeypatch): monkeypatch.delenv('MARTA_API_KEY', raising=False) with pytest.raises(APIKeyError, match='API Key is missing'): - MARTA().get_buses() + MARTA().get_trains() @pytest.mark.parametrize('status_code', [401, 403]) diff --git a/tests/test_vehicles.py b/tests/test_vehicles.py index ba3f77d..91fddf8 100644 --- a/tests/test_vehicles.py +++ b/tests/test_vehicles.py @@ -1,5 +1,7 @@ import datetime +from google.transit import gtfs_realtime_pb2 + from marta.vehicles import Bus, Train @@ -19,6 +21,36 @@ def test_bus(bus_record): assert bus.vehicle == "1469" +def test_bus_from_gtfs_vehicle_position(): + vehicle_position = gtfs_realtime_pb2.VehiclePosition() + vehicle_position.trip.route_id = '1' + vehicle_position.trip.trip_id = 'trip-1' + vehicle_position.trip.direction_id = 0 + vehicle_position.vehicle.id = '1469' + vehicle_position.vehicle.label = 'Bus 1469' + vehicle_position.position.latitude = 33.771677 + vehicle_position.position.longitude = -84.386794 + vehicle_position.position.bearing = 180 + vehicle_position.position.speed = 10 + vehicle_position.stop_id = '907473' + vehicle_position.timestamp = 1_725_000_000 + + bus = Bus.from_gtfs(vehicle_position) + + assert bus.route == 1 + assert bus.trip_id == 'trip-1' + assert bus.direction is None + assert bus.direction_id == 0 + assert bus.vehicle == '1469' + assert bus.stop_id == '907473' + assert bus.adherence is None + assert bus.timepoint is None + assert bus.latitude == vehicle_position.position.latitude + assert bus.longitude == vehicle_position.position.longitude + assert bus.last_updated.tzinfo is datetime.timezone.utc + assert bus.raw_data['vehicle']['label'] == 'Bus 1469' + + def test_train(train_record): train = Train(train_record) assert train.destination == "Doraville" From 1f8f115609779f9d9c6054909f94a4308f7c4f90 Mon Sep 17 00:00:00 2001 From: Lex Alexander Date: Sat, 5 Sep 2026 13:30:14 -0400 Subject: [PATCH 3/5] Add bus and train services --- README.md | 45 +++++ marta/__init__.py | 5 +- marta/_shared.py | 166 ++++++++++++++++++ marta/api.py | 292 +++++--------------------------- marta/bus_service.py | 5 + marta/entities/__init__.py | 7 + marta/entities/bus.py | 90 ++++++++++ marta/entities/train.py | 17 ++ marta/entities/vehicle.py | 4 + marta/services/__init__.py | 6 + marta/services/bus_service.py | 67 ++++++++ marta/services/train_service.py | 54 ++++++ marta/train_service.py | 5 + marta/vehicles.py | 110 +----------- setup.py | 2 +- 15 files changed, 517 insertions(+), 358 deletions(-) create mode 100644 marta/_shared.py create mode 100644 marta/bus_service.py create mode 100644 marta/entities/__init__.py create mode 100644 marta/entities/bus.py create mode 100644 marta/entities/train.py create mode 100644 marta/entities/vehicle.py create mode 100644 marta/services/__init__.py create mode 100644 marta/services/bus_service.py create mode 100644 marta/services/train_service.py create mode 100644 marta/train_service.py diff --git a/README.md b/README.md index f4d9fed..8cd30e4 100644 --- a/README.md +++ b/README.md @@ -98,6 +98,51 @@ trains = marta.get_trains(line='blue', station='Five Points Station', destinatio trains = marta.get_trains(line='gold', station='Five Points Station', direction='north') ``` +## Package layout + +```text +marta/ + api.py # Client composing the services + services/ + bus_service.py # Bus queries + train_service.py # Train queries + entities/ + vehicle.py # Shared vehicle base + bus.py # Bus data object + train.py # Train data object +``` + +Import services from `marta.services` and data objects from `marta.entities`. +Existing imports from `marta`, `marta.vehicles`, `marta.bus_service`, and +`marta.train_service` remain supported through re-exports. + +## Services + +`MARTA` composes a `MartaBusService` at `marta.buses` and a `MartaTrainService` +at `marta.trains`. The existing top-level query methods delegate to these services +and keep their existing arguments and return types. + +```python +from marta import MARTA, MartaBusService, MartaTrainService + +marta = MARTA() +buses = marta.buses.get_buses(route=1) +positions = marta.buses.get_vehicle_positions() +trip_updates = marta.buses.get_trip_updates() +trains = marta.trains.get_trains(line='blue') + +# Services can also be used independently. +buses = MartaBusService().get_buses(route=1) +trains = MartaTrainService(api_key='your-key').get_trains(station='Midtown Station') +``` + +Bus queries do not require an API key. The train service reads `MARTA_API_KEY` +when constructed unless a key is supplied, and accepts an `api_key` override on +individual calls. Both services share the library's existing in-memory cache, +request timeouts, and response validation. Constructing a service makes no requests. +Bus trip updates remain raw protobuf; this refactor does not add Python bus arrival +or departure projections. + ## Results The library includes very basic objects to represent Trains and Buses. Below are dictionary representations of each. diff --git a/marta/__init__.py b/marta/__init__.py index 56d70f4..01f1769 100644 --- a/marta/__init__.py +++ b/marta/__init__.py @@ -1 +1,4 @@ -from .api import MARTA \ No newline at end of file +from .api import MARTA +from .services import MartaBusService, MartaTrainService + +__all__ = ['MARTA', 'MartaBusService', 'MartaTrainService'] diff --git a/marta/_shared.py b/marta/_shared.py new file mode 100644 index 0000000..be02d1e --- /dev/null +++ b/marta/_shared.py @@ -0,0 +1,166 @@ +"""Shared transport, authentication, and filtering for MARTA services.""" + +import json +from functools import wraps +from os import getenv +from typing import List, Union +from urllib.parse import urlparse + +import requests_cache +from google.transit import gtfs_realtime_pb2 + +from .exceptions import APIKeyError, InvalidDirectionError + +_CACHE_EXPIRE = int(getenv('MARTA_CACHE_EXPIRE', 30)) +_BASE_URL = 'https://developer.itsmarta.com' +_TRAIN_URL = ( + 'https://developerservices.itsmarta.com:18096/itsmarta/' + 'railrealtimearrivals/developerservices/traindata' +) +_BUS_VEHICLE_POSITIONS_URL = ( + 'https://gtfs-rt.itsmarta.com/TMGTFSRealTimeWebService/' + 'vehicle/vehiclepositions.pb' +) +_BUS_TRIP_UPDATES_URL = ( + 'https://gtfs-rt.itsmarta.com/TMGTFSRealTimeWebService/' + 'tripupdate/tripupdates.pb' +) +_REQUEST_TIMEOUT = (3.05, 10) +_MAX_RESPONSE_BYTES = 5 * 1024 * 1024 +_JSON_CONTENT_TYPES = {'application/json', 'text/json'} +_PROTOBUF_CONTENT_TYPES = { + 'application/octet-stream', + 'application/protocol-buffer', + 'application/protobuf', + 'application/x-protobuf', +} + +# Keep caching private to this library and in memory. Importing marta must not +# patch requests globally or persist request URLs containing credentials. +CACHE = requests_cache.CachedSession( + backend='memory', + expire_after=_CACHE_EXPIRE, + ignored_parameters=['apikey', 'apiKey'], +) + +def require_api_key(func): + """ + Decorator to ensure an API key is present + """ + @wraps(func) + def with_key(self, *args, **kwargs): + if not kwargs.get('api_key'): + if not self._api_key: + raise APIKeyError() + kwargs['api_key'] = self._api_key + return func(self, *args, **kwargs) + + return with_key + +def get_bus_direction(user_direction) -> str: + direction = user_direction.lower() + if direction.startswith('n'): + return 'Northbound' + elif direction.startswith('s'): + return 'Southbound' + elif direction.startswith('e'): + return 'Eastbound' + elif direction.startswith('w'): + return 'Westbound' + else: + return None + +def get_train_direction(user_direction) -> str: + train_direction = user_direction.lower() + if train_direction.startswith('n'): + return 'N' + elif train_direction.startswith('s'): + return 'S' + elif train_direction.startswith('e'): + return 'E' + elif train_direction.startswith('w'): + return 'W' + else: + return None + +def _read_response(response, allowed_content_types) -> bytes: + response.raise_for_status() + + content_type = response.headers.get('Content-Type', '').split(';', 1)[0].lower() + if content_type and content_type not in allowed_content_types: + raise ValueError(f'Unexpected MARTA API content type: {content_type}') + + content_length = response.headers.get('Content-Length') + if content_length and int(content_length) > _MAX_RESPONSE_BYTES: + raise ValueError('MARTA API response exceeds the maximum allowed size') + + chunks = [] + received = 0 + for chunk in response.iter_content(chunk_size=64 * 1024): + received += len(chunk) + if received > _MAX_RESPONSE_BYTES: + raise ValueError('MARTA API response exceeds the maximum allowed size') + chunks.append(chunk) + return b''.join(chunks) + + +def _get_json_data( + url: str, + api_key: str, + api_key_parameter: str, +) -> List[dict]: + response = CACHE.get( + url, + params={api_key_parameter: api_key}, + timeout=_REQUEST_TIMEOUT, + stream=True, + ) + if response.status_code in (401, 403): + raise APIKeyError('Your MARTA API key was rejected.') + content = _read_response(response, _JSON_CONTENT_TYPES) + data = json.loads(content.decode(response.encoding or 'utf-8')) + if not isinstance(data, list) or any(not isinstance(item, dict) for item in data): + raise ValueError('MARTA API response must be a list of objects') + return data + + +def _get_data(endpoint: str, api_key: str) -> List[dict]: + """Fetch a legacy MARTA JSON endpoint.""" + return _get_json_data( + url=f'{_BASE_URL}{endpoint}', + api_key=api_key, + api_key_parameter='apikey', + ) + + +def _get_gtfs_realtime_feed(url: str) -> gtfs_realtime_pb2.FeedMessage: + response = CACHE.get( + url, + timeout=_REQUEST_TIMEOUT, + stream=True, + allow_redirects=False, + ) + if response.is_redirect or response.is_permanent_redirect: + redirect_url = response.headers.get('Location', '') + if urlparse(redirect_url).scheme != 'https': + raise ValueError('MARTA API attempted an insecure redirect') + raise ValueError('MARTA API returned an unexpected redirect') + content = _read_response(response, _PROTOBUF_CONTENT_TYPES) + feed = gtfs_realtime_pb2.FeedMessage() + feed.ParseFromString(content) + return feed + +def _filter_response(response: dict, filters: dict) -> List[dict]: + valid_items = [] + for item in response: + valid = True + for filter_key, filter_value in filters.items(): + if filter_value: # ignore if the filter value doesn't exist + if not item.get(filter_key): # don't penalize if item doesn't have a filter_key + pass + elif str(item[filter_key]).lower() != str(filter_value).lower(): + # lower all values to avoid case issues + valid = False + if valid: + valid_items.append(item) + return valid_items diff --git a/marta/api.py b/marta/api.py index 0128a4c..c72c17a 100644 --- a/marta/api.py +++ b/marta/api.py @@ -1,194 +1,40 @@ -import json -from functools import wraps +"""Public MARTA client composing bus and train services.""" + from os import getenv -from typing import List, Union -from urllib.parse import urlparse +from typing import List -import requests_cache from google.transit import gtfs_realtime_pb2 -from .exceptions import APIKeyError, InvalidDirectionError -from .vehicles import Bus, Train - -_CACHE_EXPIRE = int(getenv('MARTA_CACHE_EXPIRE', 30)) -_BASE_URL = 'https://developer.itsmarta.com' -_TRAIN_URL = ( - 'https://developerservices.itsmarta.com:18096/itsmarta/' - 'railrealtimearrivals/developerservices/traindata' -) -_BUS_VEHICLE_POSITIONS_URL = ( - 'https://gtfs-rt.itsmarta.com/TMGTFSRealTimeWebService/' - 'vehicle/vehiclepositions.pb' -) -_BUS_TRIP_UPDATES_URL = ( - 'https://gtfs-rt.itsmarta.com/TMGTFSRealTimeWebService/' - 'tripupdate/tripupdates.pb' +# Retain the existing module-level helpers and shared cache for compatibility. +from ._shared import ( + CACHE, + _BASE_URL, + _BUS_TRIP_UPDATES_URL, + _BUS_VEHICLE_POSITIONS_URL, + _CACHE_EXPIRE, + _JSON_CONTENT_TYPES, + _MAX_RESPONSE_BYTES, + _PROTOBUF_CONTENT_TYPES, + _REQUEST_TIMEOUT, + _TRAIN_URL, + _filter_response, + _get_data, + _get_gtfs_realtime_feed, + _get_json_data, + _read_response, + get_bus_direction, + get_train_direction, + require_api_key, ) -_REQUEST_TIMEOUT = (3.05, 10) -_MAX_RESPONSE_BYTES = 5 * 1024 * 1024 -_JSON_CONTENT_TYPES = {'application/json', 'text/json'} -_PROTOBUF_CONTENT_TYPES = { - 'application/octet-stream', - 'application/protocol-buffer', - 'application/protobuf', - 'application/x-protobuf', -} - -# Keep caching private to this library and in memory. Importing marta must not -# patch requests globally or persist request URLs containing credentials. -CACHE = requests_cache.CachedSession( - backend='memory', - expire_after=_CACHE_EXPIRE, - ignored_parameters=['apikey', 'apiKey'], -) - -def require_api_key(func): - """ - Decorator to ensure an API key is present - """ - @wraps(func) - def with_key(self, *args, **kwargs): - if not kwargs.get('api_key'): - if not self._api_key: - raise APIKeyError() - kwargs['api_key'] = self._api_key - return func(self, *args, **kwargs) - - return with_key - -def get_bus_direction(user_direction) -> str: - direction = user_direction.lower() - if direction.startswith('n'): - return 'Northbound' - elif direction.startswith('s'): - return 'Southbound' - elif direction.startswith('e'): - return 'Eastbound' - elif direction.startswith('w'): - return 'Westbound' - else: - return None - -def get_train_direction(user_direction) -> str: - train_direction = user_direction.lower() - if train_direction.startswith('n'): - return 'N' - elif train_direction.startswith('s'): - return 'S' - elif train_direction.startswith('e'): - return 'E' - elif train_direction.startswith('w'): - return 'W' - else: - return None - - -def _convert_direction(user_direction: str, vehicle_type: str = 'bus') -> Union[str, None]: - if not user_direction: - return None - if vehicle_type == 'bus': - bus_direction = get_bus_direction(user_direction) - if bus_direction is not None: - return bus_direction - else: - raise InvalidDirectionError(direction_provided=user_direction) - elif vehicle_type == 'train': - train_direction = get_train_direction(user_direction) - if train_direction is not None: - return train_direction - else: - raise InvalidDirectionError(direction_provided=user_direction) - else: - return user_direction - - -def _read_response(response, allowed_content_types) -> bytes: - response.raise_for_status() +from .services import MartaBusService, MartaTrainService +from .entities import Bus, Train - content_type = response.headers.get('Content-Type', '').split(';', 1)[0].lower() - if content_type and content_type not in allowed_content_types: - raise ValueError(f'Unexpected MARTA API content type: {content_type}') - - content_length = response.headers.get('Content-Length') - if content_length and int(content_length) > _MAX_RESPONSE_BYTES: - raise ValueError('MARTA API response exceeds the maximum allowed size') - - chunks = [] - received = 0 - for chunk in response.iter_content(chunk_size=64 * 1024): - received += len(chunk) - if received > _MAX_RESPONSE_BYTES: - raise ValueError('MARTA API response exceeds the maximum allowed size') - chunks.append(chunk) - return b''.join(chunks) - - -def _get_json_data( - url: str, - api_key: str, - api_key_parameter: str, -) -> List[dict]: - response = CACHE.get( - url, - params={api_key_parameter: api_key}, - timeout=_REQUEST_TIMEOUT, - stream=True, - ) - if response.status_code in (401, 403): - raise APIKeyError('Your MARTA API key was rejected.') - content = _read_response(response, _JSON_CONTENT_TYPES) - data = json.loads(content.decode(response.encoding or 'utf-8')) - if not isinstance(data, list) or any(not isinstance(item, dict) for item in data): - raise ValueError('MARTA API response must be a list of objects') - return data - - -def _get_data(endpoint: str, api_key: str) -> List[dict]: - """Fetch a legacy MARTA JSON endpoint.""" - return _get_json_data( - url=f'{_BASE_URL}{endpoint}', - api_key=api_key, - api_key_parameter='apikey', - ) - - -def _get_gtfs_realtime_feed(url: str) -> gtfs_realtime_pb2.FeedMessage: - response = CACHE.get( - url, - timeout=_REQUEST_TIMEOUT, - stream=True, - allow_redirects=False, - ) - if response.is_redirect or response.is_permanent_redirect: - redirect_url = response.headers.get('Location', '') - if urlparse(redirect_url).scheme != 'https': - raise ValueError('MARTA API attempted an insecure redirect') - raise ValueError('MARTA API returned an unexpected redirect') - content = _read_response(response, _PROTOBUF_CONTENT_TYPES) - feed = gtfs_realtime_pb2.FeedMessage() - feed.ParseFromString(content) - return feed - -def _filter_response(response: dict, filters: dict) -> List[dict]: - valid_items = [] - for item in response: - valid = True - for filter_key, filter_value in filters.items(): - if filter_value: # ignore if the filter value doesn't exist - if not item.get(filter_key): # don't penalize if item doesn't have a filter_key - pass - elif str(item[filter_key]).lower() != str(filter_value).lower(): - # lower all values to avoid case issues - valid = False - if valid: - valid_items.append(item) - return valid_items class MARTA: def __init__(self, api_key: str = None): - self._api_key = api_key - if not api_key: - self._api_key = getenv('MARTA_API_KEY') + self._api_key = api_key or getenv('MARTA_API_KEY') + self.buses = MartaBusService() + self.trains = MartaTrainService(api_key=self._api_key) @require_api_key def get_trains(self, @@ -197,82 +43,28 @@ def get_trains(self, destination: str = None, direction: str = None, api_key: str = None) -> List[Train]: - """ - Query API for train information - - :param line: Train line identifier filter (red, gold, green, or blue) - :type line: str, optional - :param station: train station filter - :type station: str, optional - :param destination: destination filter - :type destination: str, optional - :param direction: Direction train is heading (N, S, E, or W) - :param api_key: API key to override environment variable - :type api_key: str, optional - :return: list of Train objects - :rtype: List[Train] - """ - data = _get_json_data( - url=_TRAIN_URL, - api_key=api_key, - api_key_parameter='apiKey', + """Delegate train queries to :attr:`trains`, preserving existing filters.""" + return self.trains.get_trains( + line=line, station=station, destination=destination, + direction=direction, api_key=api_key, ) - filters = { - 'LINE': line, - 'DIRECTION': _convert_direction(user_direction=direction, vehicle_type='train'), - 'STATION': station, - 'DESTINATION': destination - } - matching_data = _filter_response(response=data, filters=filters) - return [Train(t) for t in matching_data] def get_buses(self, route: int = None, stop_id: int = None, vehicle_id: int = None, time_point: str = None, - direction: str = None - ) -> List[Bus]: - """ - Query the GTFS-Realtime vehicle feed for bus information. - :param route: route number - :type route: int, optional - :param stop_id: Bus stop ID - :type stop_id: int, optional - :param vehicle_id: Bus ID - :type vehicle_id: int, optional - :param time_point: Deprecated; unavailable in GTFS-Realtime - :type time_point: str, optional - :param direction: Deprecated; cardinal direction is unavailable in GTFS-Realtime - :type direction: str, optional - :param api_key: Deprecated; the GTFS-Realtime feed does not require a key - :type api_key: str, optional - :return: list of Bus objects - """ - - if time_point is not None or direction is not None: - raise ValueError( - 'time_point and direction filters are unavailable in ' - 'MARTA GTFS-Realtime data' - ) - - feed = self.get_bus_vehicle_positions() - buses = [ - Bus.from_gtfs(entity.vehicle) - for entity in feed.entity - if entity.HasField('vehicle') - ] - return [ - bus for bus in buses - if (route is None or str(bus.route) == str(route)) - and (stop_id is None or str(bus.stop_id) == str(stop_id)) - and (vehicle_id is None or str(bus.vehicle) == str(vehicle_id)) - ] + direction: str = None) -> List[Bus]: + """Delegate bus queries to :attr:`buses`, returning Bus objects.""" + return self.buses.get_buses( + route=route, stop_id=stop_id, vehicle_id=vehicle_id, + time_point=time_point, direction=direction, + ) def get_bus_vehicle_positions(self) -> gtfs_realtime_pb2.FeedMessage: - """Return MARTA's GTFS-Realtime bus vehicle-position feed.""" - return _get_gtfs_realtime_feed(_BUS_VEHICLE_POSITIONS_URL) + """Return the bus service's raw GTFS-Realtime vehicle-position feed.""" + return self.buses.get_vehicle_positions() def get_bus_trip_updates(self) -> gtfs_realtime_pb2.FeedMessage: - """Return MARTA's GTFS-Realtime bus trip-update feed.""" - return _get_gtfs_realtime_feed(_BUS_TRIP_UPDATES_URL) + """Return the bus service's raw GTFS-Realtime trip-update feed.""" + return self.buses.get_trip_updates() diff --git a/marta/bus_service.py b/marta/bus_service.py new file mode 100644 index 0000000..c8e4d25 --- /dev/null +++ b/marta/bus_service.py @@ -0,0 +1,5 @@ +"""Compatibility import; use marta.services instead.""" + +from .services.bus_service import MartaBusService + +__all__ = ['MartaBusService'] diff --git a/marta/entities/__init__.py b/marta/entities/__init__.py new file mode 100644 index 0000000..4d1b832 --- /dev/null +++ b/marta/entities/__init__.py @@ -0,0 +1,7 @@ +"""Vehicle data objects returned by MARTA query services.""" + +from .bus import Bus +from .train import Train +from .vehicle import Vehicle + +__all__ = ['Bus', 'Train', 'Vehicle'] diff --git a/marta/entities/bus.py b/marta/entities/bus.py new file mode 100644 index 0000000..e7eab8c --- /dev/null +++ b/marta/entities/bus.py @@ -0,0 +1,90 @@ +from datetime import datetime, timezone + +from .vehicle import Vehicle + + +class Bus(Vehicle): + def __init__(self, record): + self.raw_data = record + self.adherence = record.get('ADHERENCE') + self.block_id = record.get('BLOCKID') + self.block_abbr = record.get('BLOCK_ABBR') + self.direction = record.get('DIRECTION') + self.latitude = record.get('LATITUDE') + self.longitude = record.get('LONGITUDE') + self.last_updated = datetime.strptime(record.get('MSGTIME'), '%m/%d/%Y %H:%M:%S %p') + self.route = int(record.get('ROUTE')) + self.stop_id = record.get('STOPID') + self.timepoint = record.get('TIMEPOINT') + self.trip_id = record.get('TRIPID') + self.vehicle = record.get('VEHICLE') + + @classmethod + def from_gtfs(cls, vehicle_position): + """Build a compatibility Bus from a GTFS-Realtime VehiclePosition.""" + bus = cls.__new__(cls) + trip = vehicle_position.trip + vehicle = vehicle_position.vehicle + position = vehicle_position.position + + bus.adherence = None + bus.block_id = None + bus.block_abbr = None + bus.direction = None + bus.direction_id = ( + trip.direction_id if trip.HasField('direction_id') else None + ) + bus.latitude = position.latitude + bus.longitude = position.longitude + bus.bearing = position.bearing if position.HasField('bearing') else None + bus.speed = position.speed if position.HasField('speed') else None + bus.last_updated = ( + datetime.fromtimestamp(vehicle_position.timestamp, timezone.utc) + if vehicle_position.HasField('timestamp') + else None + ) + bus.route = _coerce_route(trip.route_id) + bus.stop_id = vehicle_position.stop_id or None + bus.timepoint = None + bus.trip_id = trip.trip_id or None + bus.vehicle = vehicle.id or None + bus.current_status = ( + vehicle_position.current_status + if vehicle_position.HasField('current_status') + else None + ) + bus.raw_data = { + 'trip': { + 'trip_id': bus.trip_id, + 'route_id': trip.route_id or None, + 'direction_id': bus.direction_id, + }, + 'vehicle': { + 'id': bus.vehicle, + 'label': vehicle.label or None, + 'license_plate': vehicle.license_plate or None, + }, + 'position': { + 'latitude': bus.latitude, + 'longitude': bus.longitude, + 'bearing': bus.bearing, + 'speed': bus.speed, + }, + 'stop_id': bus.stop_id, + 'timestamp': ( + vehicle_position.timestamp + if vehicle_position.HasField('timestamp') + else None + ), + 'current_status': bus.current_status, + } + return bus + + +def _coerce_route(route_id): + if not route_id: + return None + try: + return int(route_id) + except ValueError: + return route_id diff --git a/marta/entities/train.py b/marta/entities/train.py new file mode 100644 index 0000000..c71945e --- /dev/null +++ b/marta/entities/train.py @@ -0,0 +1,17 @@ +from datetime import datetime + +from .vehicle import Vehicle + + +class Train(Vehicle): + def __init__(self, record): + self.raw_data = record + self.destination = record.get('DESTINATION') + self.direction = record.get('DIRECTION') + self.last_updated = datetime.strptime(record.get('EVENT_TIME'), '%m/%d/%Y %H:%M:%S %p') + self.line = record.get('LINE') + self.next_arrival = datetime.strptime(record.get('NEXT_ARR'), '%H:%M:%S %p').time() + self.station = record.get('STATION') + self.train_id = record.get('TRAIN_ID') + self.waiting_seconds = record.get('WAITING_SECONDS') + self.waiting_time = record.get('WAITING_TIME') diff --git a/marta/entities/vehicle.py b/marta/entities/vehicle.py new file mode 100644 index 0000000..0dd3db2 --- /dev/null +++ b/marta/entities/vehicle.py @@ -0,0 +1,4 @@ +class Vehicle(): + """Generic Vehicle object that exists to print vehicles as dicts""" + def __str__(self): + return str(self.__dict__) diff --git a/marta/services/__init__.py b/marta/services/__init__.py new file mode 100644 index 0000000..9a78051 --- /dev/null +++ b/marta/services/__init__.py @@ -0,0 +1,6 @@ +"""MARTA bus and train query services.""" + +from .bus_service import MartaBusService +from .train_service import MartaTrainService + +__all__ = ['MartaBusService', 'MartaTrainService'] diff --git a/marta/services/bus_service.py b/marta/services/bus_service.py new file mode 100644 index 0000000..2ad66bb --- /dev/null +++ b/marta/services/bus_service.py @@ -0,0 +1,67 @@ +"""Bus queries backed by MARTA's GTFS-Realtime feeds.""" + +from typing import List + +from google.transit import gtfs_realtime_pb2 + +from .._shared import ( + _BUS_TRIP_UPDATES_URL, + _BUS_VEHICLE_POSITIONS_URL, + _get_gtfs_realtime_feed, +) +from ..entities import Bus + + +class MartaBusService: + """Query bus positions and trip updates; no API key is required.""" + + def get_buses(self, + route: int = None, + stop_id: int = None, + vehicle_id: int = None, + time_point: str = None, + direction: str = None + ) -> List[Bus]: + """ + Query the GTFS-Realtime vehicle feed for bus information. + :param route: route number + :type route: int, optional + :param stop_id: Bus stop ID + :type stop_id: int, optional + :param vehicle_id: Bus ID + :type vehicle_id: int, optional + :param time_point: Deprecated; unavailable in GTFS-Realtime + :type time_point: str, optional + :param direction: Deprecated; cardinal direction is unavailable in GTFS-Realtime + :type direction: str, optional + :param api_key: Deprecated; the GTFS-Realtime feed does not require a key + :type api_key: str, optional + :return: list of Bus objects + """ + + if time_point is not None or direction is not None: + raise ValueError( + 'time_point and direction filters are unavailable in ' + 'MARTA GTFS-Realtime data' + ) + + feed = self.get_vehicle_positions() + buses = [ + Bus.from_gtfs(entity.vehicle) + for entity in feed.entity + if entity.HasField('vehicle') + ] + return [ + bus for bus in buses + if (route is None or str(bus.route) == str(route)) + and (stop_id is None or str(bus.stop_id) == str(stop_id)) + and (vehicle_id is None or str(bus.vehicle) == str(vehicle_id)) + ] + + def get_vehicle_positions(self) -> gtfs_realtime_pb2.FeedMessage: + """Return MARTA's GTFS-Realtime bus vehicle-position feed.""" + return _get_gtfs_realtime_feed(_BUS_VEHICLE_POSITIONS_URL) + + def get_trip_updates(self) -> gtfs_realtime_pb2.FeedMessage: + """Return MARTA's GTFS-Realtime bus trip-update feed.""" + return _get_gtfs_realtime_feed(_BUS_TRIP_UPDATES_URL) diff --git a/marta/services/train_service.py b/marta/services/train_service.py new file mode 100644 index 0000000..a4fc4f1 --- /dev/null +++ b/marta/services/train_service.py @@ -0,0 +1,54 @@ +"""Train queries backed by MARTA's realtime rail API.""" + +from os import getenv +from typing import List + +from .._shared import _TRAIN_URL, get_train_direction, _filter_response, _get_json_data, require_api_key +from ..entities import Train +from ..exceptions import InvalidDirectionError + + +class MartaTrainService: + """Query train arrivals with a configured or per-call API key.""" + + def __init__(self, api_key: str = None): + self._api_key = api_key or getenv('MARTA_API_KEY') + + @require_api_key + def get_trains(self, + line: str = None, + station: str = None, + destination: str = None, + direction: str = None, + api_key: str = None) -> List[Train]: + """ + Query API for train information + + :param line: Train line identifier filter (red, gold, green, or blue) + :type line: str, optional + :param station: train station filter + :type station: str, optional + :param destination: destination filter + :type destination: str, optional + :param direction: Direction train is heading (N, S, E, or W) + :param api_key: API key to override environment variable + :type api_key: str, optional + :return: list of Train objects + :rtype: List[Train] + """ + data = _get_json_data( + url=_TRAIN_URL, + api_key=api_key, + api_key_parameter='apiKey', + ) + normalized_direction = get_train_direction(direction) if direction else None + if direction and normalized_direction is None: + raise InvalidDirectionError(direction_provided=direction) + filters = { + 'LINE': line, + 'DIRECTION': normalized_direction, + 'STATION': station, + 'DESTINATION': destination + } + matching_data = _filter_response(response=data, filters=filters) + return [Train(t) for t in matching_data] diff --git a/marta/train_service.py b/marta/train_service.py new file mode 100644 index 0000000..3ebd940 --- /dev/null +++ b/marta/train_service.py @@ -0,0 +1,5 @@ +"""Compatibility import; use marta.services instead.""" + +from .services.train_service import MartaTrainService + +__all__ = ['MartaTrainService'] diff --git a/marta/vehicles.py b/marta/vehicles.py index 02cb6da..e9f2a06 100644 --- a/marta/vehicles.py +++ b/marta/vehicles.py @@ -1,108 +1,6 @@ -from datetime import datetime, timezone +"""Compatibility imports; vehicle entities now live in marta.entities.""" +from .entities import Bus, Train, Vehicle +from .entities.bus import _coerce_route -class Vehicle(): - """Generic Vehicle object that exists to print vehicles as dicts""" - def __str__(self): - return str(self.__dict__) - - -class Bus(Vehicle): - def __init__(self, record): - self.raw_data = record - self.adherence = record.get('ADHERENCE') - self.block_id = record.get('BLOCKID') - self.block_abbr = record.get('BLOCK_ABBR') - self.direction = record.get('DIRECTION') - self.latitude = record.get('LATITUDE') - self.longitude = record.get('LONGITUDE') - self.last_updated = datetime.strptime(record.get('MSGTIME'), '%m/%d/%Y %H:%M:%S %p') - self.route = int(record.get('ROUTE')) - self.stop_id = record.get('STOPID') - self.timepoint = record.get('TIMEPOINT') - self.trip_id = record.get('TRIPID') - self.vehicle = record.get('VEHICLE') - - @classmethod - def from_gtfs(cls, vehicle_position): - """Build a compatibility Bus from a GTFS-Realtime VehiclePosition.""" - bus = cls.__new__(cls) - trip = vehicle_position.trip - vehicle = vehicle_position.vehicle - position = vehicle_position.position - - bus.adherence = None - bus.block_id = None - bus.block_abbr = None - bus.direction = None - bus.direction_id = ( - trip.direction_id if trip.HasField('direction_id') else None - ) - bus.latitude = position.latitude - bus.longitude = position.longitude - bus.bearing = position.bearing if position.HasField('bearing') else None - bus.speed = position.speed if position.HasField('speed') else None - bus.last_updated = ( - datetime.fromtimestamp(vehicle_position.timestamp, timezone.utc) - if vehicle_position.HasField('timestamp') - else None - ) - bus.route = _coerce_route(trip.route_id) - bus.stop_id = vehicle_position.stop_id or None - bus.timepoint = None - bus.trip_id = trip.trip_id or None - bus.vehicle = vehicle.id or None - bus.current_status = ( - vehicle_position.current_status - if vehicle_position.HasField('current_status') - else None - ) - bus.raw_data = { - 'trip': { - 'trip_id': bus.trip_id, - 'route_id': trip.route_id or None, - 'direction_id': bus.direction_id, - }, - 'vehicle': { - 'id': bus.vehicle, - 'label': vehicle.label or None, - 'license_plate': vehicle.license_plate or None, - }, - 'position': { - 'latitude': bus.latitude, - 'longitude': bus.longitude, - 'bearing': bus.bearing, - 'speed': bus.speed, - }, - 'stop_id': bus.stop_id, - 'timestamp': ( - vehicle_position.timestamp - if vehicle_position.HasField('timestamp') - else None - ), - 'current_status': bus.current_status, - } - return bus - - -def _coerce_route(route_id): - if not route_id: - return None - try: - return int(route_id) - except ValueError: - return route_id - - -class Train(Vehicle): - def __init__(self, record): - self.raw_data = record - self.destination = record.get('DESTINATION') - self.direction = record.get('DIRECTION') - self.last_updated = datetime.strptime(record.get('EVENT_TIME'), '%m/%d/%Y %H:%M:%S %p') - self.line = record.get('LINE') - self.next_arrival = datetime.strptime(record.get('NEXT_ARR'), '%H:%M:%S %p').time() - self.station = record.get('STATION') - self.train_id = record.get('TRAIN_ID') - self.waiting_seconds = record.get('WAITING_SECONDS') - self.waiting_time = record.get('WAITING_TIME') +__all__ = ['Bus', 'Train', 'Vehicle'] diff --git a/setup.py b/setup.py index 36c532a..612e20b 100644 --- a/setup.py +++ b/setup.py @@ -4,7 +4,7 @@ name='marta', description='Python library for accessing MARTA real-time API', url='https://www.itsmarta.com/app-developer-resources.aspx', - packages=['marta'], + packages=['marta', 'marta.services', 'marta.entities'], python_requires='>=3.8', install_requires=[ 'gtfs-realtime-bindings>=2.2,<3', From 6f1ddb0b0434f43d7708bb38d2ff04750d21d783 Mon Sep 17 00:00:00 2001 From: Lex Alexander Date: Mon, 7 Sep 2026 18:41:42 -0400 Subject: [PATCH 4/5] Explicitly check for API key in APIs that require it --- marta/_shared.py | 17 ++---- marta/api.py | 20 +++---- marta/services/bus_service.py | 7 ++- marta/services/train_service.py | 10 ++-- tests/test_api.py | 18 +++--- tests/test_package_layout.py | 16 ++++++ tests/test_services.py | 99 +++++++++++++++++++++++++++++++++ 7 files changed, 149 insertions(+), 38 deletions(-) create mode 100644 tests/test_package_layout.py create mode 100644 tests/test_services.py diff --git a/marta/_shared.py b/marta/_shared.py index be02d1e..e729c17 100644 --- a/marta/_shared.py +++ b/marta/_shared.py @@ -6,6 +6,7 @@ from typing import List, Union from urllib.parse import urlparse +from openai import api_key import requests_cache from google.transit import gtfs_realtime_pb2 @@ -43,19 +44,9 @@ ignored_parameters=['apikey', 'apiKey'], ) -def require_api_key(func): - """ - Decorator to ensure an API key is present - """ - @wraps(func) - def with_key(self, *args, **kwargs): - if not kwargs.get('api_key'): - if not self._api_key: - raise APIKeyError() - kwargs['api_key'] = self._api_key - return func(self, *args, **kwargs) - - return with_key +def require_api_key(api_key): + if api_key is None: + raise APIKeyError("API key is required but not provided.") def get_bus_direction(user_direction) -> str: direction = user_direction.lower() diff --git a/marta/api.py b/marta/api.py index c72c17a..8f0697f 100644 --- a/marta/api.py +++ b/marta/api.py @@ -4,7 +4,7 @@ from typing import List from google.transit import gtfs_realtime_pb2 - +from .exceptions import APIKeyError # Retain the existing module-level helpers and shared cache for compatibility. from ._shared import ( CACHE, @@ -32,23 +32,23 @@ class MARTA: def __init__(self, api_key: str = None): - self._api_key = api_key or getenv('MARTA_API_KEY') - self.buses = MartaBusService() - self.trains = MartaTrainService(api_key=self._api_key) + self.api_key = api_key + self.buses = MartaBusService(api_key=self.api_key) + self.trains = MartaTrainService(api_key=self.api_key) - @require_api_key def get_trains(self, line: str = None, station: str = None, destination: str = None, - direction: str = None, - api_key: str = None) -> List[Train]: + direction: str = None) -> List[Train]: """Delegate train queries to :attr:`trains`, preserving existing filters.""" return self.trains.get_trains( - line=line, station=station, destination=destination, - direction=direction, api_key=api_key, + line=line, station=station, + destination=destination, + direction=direction ) + def get_buses(self, route: int = None, stop_id: int = None, @@ -56,11 +56,11 @@ def get_buses(self, time_point: str = None, direction: str = None) -> List[Bus]: """Delegate bus queries to :attr:`buses`, returning Bus objects.""" + return self.buses.get_buses( route=route, stop_id=stop_id, vehicle_id=vehicle_id, time_point=time_point, direction=direction, ) - def get_bus_vehicle_positions(self) -> gtfs_realtime_pb2.FeedMessage: """Return the bus service's raw GTFS-Realtime vehicle-position feed.""" return self.buses.get_vehicle_positions() diff --git a/marta/services/bus_service.py b/marta/services/bus_service.py index 2ad66bb..4f23345 100644 --- a/marta/services/bus_service.py +++ b/marta/services/bus_service.py @@ -8,6 +8,7 @@ _BUS_TRIP_UPDATES_URL, _BUS_VEHICLE_POSITIONS_URL, _get_gtfs_realtime_feed, + require_api_key, ) from ..entities import Bus @@ -15,6 +16,9 @@ class MartaBusService: """Query bus positions and trip updates; no API key is required.""" + def __init__(self, api_key: str = None): + self.api_key = api_key + def get_buses(self, route: int = None, stop_id: int = None, @@ -38,7 +42,6 @@ def get_buses(self, :type api_key: str, optional :return: list of Bus objects """ - if time_point is not None or direction is not None: raise ValueError( 'time_point and direction filters are unavailable in ' @@ -60,8 +63,10 @@ def get_buses(self, def get_vehicle_positions(self) -> gtfs_realtime_pb2.FeedMessage: """Return MARTA's GTFS-Realtime bus vehicle-position feed.""" + require_api_key(self.api_key) return _get_gtfs_realtime_feed(_BUS_VEHICLE_POSITIONS_URL) def get_trip_updates(self) -> gtfs_realtime_pb2.FeedMessage: """Return MARTA's GTFS-Realtime bus trip-update feed.""" + require_api_key(self.api_key) return _get_gtfs_realtime_feed(_BUS_TRIP_UPDATES_URL) diff --git a/marta/services/train_service.py b/marta/services/train_service.py index a4fc4f1..708e112 100644 --- a/marta/services/train_service.py +++ b/marta/services/train_service.py @@ -12,15 +12,15 @@ class MartaTrainService: """Query train arrivals with a configured or per-call API key.""" def __init__(self, api_key: str = None): - self._api_key = api_key or getenv('MARTA_API_KEY') - - @require_api_key + self.api_key = api_key + require_api_key(api_key) + def get_trains(self, line: str = None, station: str = None, destination: str = None, direction: str = None, - api_key: str = None) -> List[Train]: + ) -> List[Train]: """ Query API for train information @@ -38,7 +38,7 @@ def get_trains(self, """ data = _get_json_data( url=_TRAIN_URL, - api_key=api_key, + api_key=self.api_key, api_key_parameter='apiKey', ) normalized_direction = get_train_direction(direction) if direction else None diff --git a/tests/test_api.py b/tests/test_api.py index d826f52..b6978bd 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -129,7 +129,7 @@ def test_get_buses_uses_gtfs_vehicle_feed(mock_response): ) ) - buses = MARTA().get_buses() + buses = MARTA(api_key='abc123').get_buses() assert len(buses) == 2 assert all(isinstance(bus, Bus) for bus in buses) @@ -152,7 +152,7 @@ def test_get_buses_filters_gtfs_feed_locally(mock_response): ) ) - buses = MARTA().get_buses(route=1, stop_id=907473, vehicle_id=1469) + buses = MARTA(api_key='abc123').get_buses(route=1, stop_id=907473, vehicle_id=1469) assert len(buses) == 1 assert all(bus.route == 1 for bus in buses) @@ -173,7 +173,7 @@ def test_get_bus_vehicle_positions_returns_gtfs_feed(mock_response): ) ) - feed = MARTA().get_bus_vehicle_positions() + feed = MARTA(api_key='abc123').get_bus_vehicle_positions() assert isinstance(feed, gtfs_realtime_pb2.FeedMessage) assert feed.entity[0].vehicle.vehicle.id == '1469' @@ -199,7 +199,7 @@ def test_get_bus_trip_updates_returns_gtfs_feed(mock_response): ) ) - feed = MARTA().get_bus_trip_updates() + feed = MARTA(api_key='abc123').get_bus_trip_updates() assert isinstance(feed, gtfs_realtime_pb2.FeedMessage) assert feed.entity[0].trip_update.trip.trip_id == '5391405' @@ -223,7 +223,7 @@ def test_get_buses_keeps_bus_object_contract(mock_response): ) ) - buses = MARTA().get_buses() + buses = MARTA(api_key='abc123').get_buses() assert isinstance(buses, list) assert all(isinstance(bus, Bus) for bus in buses) @@ -234,7 +234,7 @@ def test_get_buses_keeps_bus_object_contract(mock_response): @pytest.mark.parametrize('filters', [{'time_point': 'Five Points'}, {'direction': 'N'}]) def test_get_buses_rejects_unsupported_legacy_filters(filters): with pytest.raises(ValueError, match='unavailable'): - MARTA().get_buses(**filters) + MARTA(api_key='abc123').get_buses(**filters) def test_rejects_invalid_protobuf(mock_response): @@ -246,7 +246,7 @@ def test_rejects_invalid_protobuf(mock_response): ) with pytest.raises(DecodeError): - MARTA().get_bus_vehicle_positions() + MARTA(api_key='abc123').get_bus_vehicle_positions() def test_rejects_insecure_gtfs_redirect(mock_response): @@ -264,13 +264,13 @@ def test_rejects_insecure_gtfs_redirect(mock_response): ) with pytest.raises(ValueError, match='insecure redirect'): - MARTA().get_bus_vehicle_positions() + MARTA(api_key='abc123').get_bus_vehicle_positions() def test_missing_api_key_is_deterministic(monkeypatch): monkeypatch.delenv('MARTA_API_KEY', raising=False) - with pytest.raises(APIKeyError, match='API Key is missing'): + with pytest.raises(APIKeyError, match='API key is required but not provided.'): MARTA().get_trains() diff --git a/tests/test_package_layout.py b/tests/test_package_layout.py new file mode 100644 index 0000000..89f4bd6 --- /dev/null +++ b/tests/test_package_layout.py @@ -0,0 +1,16 @@ +from marta import MartaBusService, MartaTrainService +from marta.bus_service import MartaBusService as LegacyBusService +from marta.train_service import MartaTrainService as LegacyTrainService +from marta.entities import Bus, Train, Vehicle +from marta.services import MartaBusService as BusService, MartaTrainService as TrainService +from marta.vehicles import Bus as LegacyBus, Train as LegacyTrain, Vehicle as LegacyVehicle + + +def test_compatibility_imports_preserve_class_identity(): + assert Bus is LegacyBus + assert Train is LegacyTrain + assert Vehicle is LegacyVehicle + assert BusService is MartaBusService is LegacyBusService + assert TrainService is MartaTrainService is LegacyTrainService + assert issubclass(Bus, Vehicle) + assert issubclass(Train, Vehicle) diff --git a/tests/test_services.py b/tests/test_services.py new file mode 100644 index 0000000..a62f5f6 --- /dev/null +++ b/tests/test_services.py @@ -0,0 +1,99 @@ +import json +from unittest.mock import Mock + +import pytest +from google.transit import gtfs_realtime_pb2 + +from marta import MARTA, MartaBusService, MartaTrainService +from marta import api, _shared +from marta.exceptions import APIKeyError, InvalidDirectionError +from marta.vehicles import Bus, Train +from test_api import FakeResponse, gtfs_vehicle_feed + + +def test_client_constructs_services_without_requests(monkeypatch): + get = Mock(side_effect=AssertionError('Construction must not fetch')) + monkeypatch.setattr(_shared.CACHE, 'get', get) + client = MARTA(api_key='secret') + assert isinstance(client.buses, MartaBusService) + assert isinstance(client.trains, MartaTrainService) + assert api.CACHE is _shared.CACHE + get.assert_not_called() + + +def test_client_delegates_bus_filters_and_raw_feeds(): + client = MARTA(api_key='abc1234') + client.buses = Mock(spec=MartaBusService) + expected = client.buses.get_buses.return_value + assert client.get_buses(route=1, stop_id=2, vehicle_id=3) is expected + client.buses.get_buses.assert_called_once_with( + route=1, stop_id=2, vehicle_id=3, time_point=None, direction=None, + ) + assert client.get_bus_vehicle_positions() is client.buses.get_vehicle_positions.return_value + assert client.get_bus_trip_updates() is client.buses.get_trip_updates.return_value + +def test_bus_service_queries_without_key_and_preserves_bus_contract(monkeypatch): + feed = gtfs_vehicle_feed( + {'route': 1, 'stop_id': 2, 'vehicle_id': 3}, + {'route': 4, 'stop_id': 2, 'vehicle_id': 5}, + ) + response = FakeResponse( + headers={'Content-Type': 'application/protocol-buffer'}, + chunks=[feed.SerializeToString()], + ) + get = Mock(return_value=response) + monkeypatch.setattr(_shared.CACHE, 'get', get) + buses = MartaBusService().get_buses(route=1, stop_id=2, vehicle_id=3) + assert len(buses) == 1 + assert isinstance(buses[0], Bus) + assert str(buses[0].route) == '1' + assert 'params' not in get.call_args.kwargs + assert get.call_args.kwargs['timeout'] == (3.05, 10) + + +@pytest.mark.parametrize('method,url', [ + ('get_vehicle_positions', _shared._BUS_VEHICLE_POSITIONS_URL), + ('get_trip_updates', _shared._BUS_TRIP_UPDATES_URL), +]) +def test_bus_service_returns_raw_feed(monkeypatch, method, url): + feed = gtfs_realtime_pb2.FeedMessage() + feed.header.gtfs_realtime_version = '2.0' + get = Mock(return_value=FakeResponse( + headers={'Content-Type': 'application/protocol-buffer'}, + chunks=[feed.SerializeToString()], + )) + monkeypatch.setattr(_shared.CACHE, 'get', get) + assert getattr(MartaBusService(), method)() == feed + assert get.call_args.args == (url,) + + +@pytest.mark.parametrize('filters', [{'time_point': 'station'}, {'direction': 'north'}]) +def test_bus_service_retains_unsupported_filter_errors(filters): + with pytest.raises(ValueError, match='unavailable'): + MartaBusService().get_buses(**filters) + + +def test_train_service_filters_and_scopes_key_override(monkeypatch, train_response): + get = Mock(return_value=FakeResponse(json.loads(train_response))) + monkeypatch.setattr(_shared.CACHE, 'get', get) + service = MartaTrainService() + trains = service.get_trains(line='blue', station='Indian Creek Station', direction='east', api_key='override') + assert len(trains) == 1 + assert isinstance(trains[0], Train) + assert trains[0].direction == 'E' + assert get.call_args.kwargs['params'] == {'apiKey': 'override'} + service.get_trains() + assert get.call_args.kwargs['params'] == {'apiKey': 'environment'} + MartaTrainService(api_key='explicit').get_trains() + assert get.call_args.kwargs['params'] == {'apiKey': 'explicit'} + + +def test_train_service_requires_key(monkeypatch): + with pytest.raises(APIKeyError): + MartaTrainService().get_trains() + + +def test_train_service_retains_invalid_direction_error(monkeypatch): + monkeypatch.setattr(_shared.CACHE, 'get', Mock(return_value=FakeResponse([]))) + with pytest.raises(InvalidDirectionError): + MartaTrainService(api_key='secret').get_trains(direction='invalid') From 1e72a638244e34f32a1b16967fcefc78c090c35f Mon Sep 17 00:00:00 2001 From: Lex Alexander Date: Tue, 8 Sep 2026 13:29:17 -0400 Subject: [PATCH 5/5] Remove accidentally openai import and update readme and remove API check decorator in favor of explict API call in methods that require it --- README.md | 14 -------------- marta/_shared.py | 2 -- tests/conftest.py | 8 ++++++++ tests/test_services.py | 40 +++++++++++++++++++++++++++++----------- 4 files changed, 37 insertions(+), 27 deletions(-) diff --git a/README.md b/README.md index 8cd30e4..b1c9350 100644 --- a/README.md +++ b/README.md @@ -33,20 +33,6 @@ On Mac/Linux: export MARTA_API_KEY= ``` -Optionally, you can also set the cache timeout (default 30 seconds): - -Windows: - -``` -set MARTA_CACHE_EXPIRE=15 -``` - -Mac/Linux: - -``` -export MARTA_CACHE_EXPIRE=15 -``` - Create a client and use `get_buses()` or `get_trains()` for the existing object interface. `get_buses()` now reads MARTA's GTFS-Realtime vehicle feed and adapts each vehicle position into a `Bus` object. diff --git a/marta/_shared.py b/marta/_shared.py index e729c17..c66fba0 100644 --- a/marta/_shared.py +++ b/marta/_shared.py @@ -5,8 +5,6 @@ from os import getenv from typing import List, Union from urllib.parse import urlparse - -from openai import api_key import requests_cache from google.transit import gtfs_realtime_pb2 diff --git a/tests/conftest.py b/tests/conftest.py index 4479935..c3b1313 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -2,6 +2,14 @@ import pytest +@pytest.fixture(autouse=True) +def block_unmocked_http(monkeypatch): + def unexpected_request(*args, **kwargs): + raise AssertionError('Unexpected HTTP request: mock marta._shared.CACHE.get') + + monkeypatch.setattr('requests.sessions.Session.request', unexpected_request) + + @pytest.fixture def train_response(): return '[{"DESTINATION":"Doraville","DIRECTION":"N","EVENT_TIME":"1\/26\/2017 9:58:38 PM","LINE":"GOLD","NEXT_ARR":"09:58:47 PM","STATION":"CHAMBLEE STATION","TRAIN_ID":"305326","WAITING_SECONDS":"-40","WAITING_TIME":"Boarding"},{"DESTINATION":"Airport","DIRECTION":"S","EVENT_TIME":"1\/26\/2017 9:59:06 PM","LINE":"GOLD","NEXT_ARR":"09:59:15 PM","STATION":"CIVIC CENTER STATION","TRAIN_ID":"303506","WAITING_SECONDS":"-12","WAITING_TIME":"Boarding"},{"DESTINATION":"Indian Creek","DIRECTION":"E","EVENT_TIME":"1\/26\/2017 9:59:16 PM","LINE":"BLUE","NEXT_ARR":"09:59:25 PM","STATION":"KING MEMORIAL STATION","TRAIN_ID":"102026","WAITING_SECONDS":"-2","WAITING_TIME":"Boarding"},{"DESTINATION":"Hamilton E Holmes","DIRECTION":"W","EVENT_TIME":"1\/26\/2017 9:58:48 PM","LINE":"BLUE","NEXT_ARR":"09:59:29 PM","STATION":"HAMILTON E HOLMES STATION","TRAIN_ID":"104206","WAITING_SECONDS":"2","WAITING_TIME":"Arriving"},{"DESTINATION":"Doraville","DIRECTION":"N","EVENT_TIME":"1\/26\/2017 9:59:16 PM","LINE":"GOLD","NEXT_ARR":"09:59:39 PM","STATION":"PEACHTREE CENTER STATION","TRAIN_ID":"306326","WAITING_SECONDS":"12","WAITING_TIME":"Arriving"},{"DESTINATION":"North Springs","DIRECTION":"N","EVENT_TIME":"1\/26\/2017 9:59:16 PM","LINE":"RED","NEXT_ARR":"10:00:04 PM","STATION":"BUCKHEAD STATION","TRAIN_ID":"410306","WAITING_SECONDS":"37","WAITING_TIME":"Arriving"},{"DESTINATION":"Hamilton E Holmes","DIRECTION":"W","EVENT_TIME":"1\/26\/2017 9:58:44 PM","LINE":"BLUE","NEXT_ARR":"10:00:22 PM","STATION":"EAST LAKE STATION","TRAIN_ID":"106206","WAITING_SECONDS":"55","WAITING_TIME":"Arriving"},{"DESTINATION":"Airport","DIRECTION":"S","EVENT_TIME":"1\/26\/2017 9:59:06 PM","LINE":"GOLD","NEXT_ARR":"10:00:27 PM","STATION":"PEACHTREE CENTER STATION","TRAIN_ID":"303506","WAITING_SECONDS":"60","WAITING_TIME":"Arriving"},{"DESTINATION":"","DIRECTION":"W","EVENT_TIME":"1\/26\/2017 9:59:00 PM","LINE":"GREEN","NEXT_ARR":"10:00:29 PM","STATION":"BANKHEAD STATION","TRAIN_ID":"201172","WAITING_SECONDS":"62","WAITING_TIME":"Arriving"},{"DESTINATION":"Doraville","DIRECTION":"N","EVENT_TIME":"1\/26\/2017 9:59:16 PM","LINE":"GOLD","NEXT_ARR":"10:01:10 PM","STATION":"CIVIC CENTER STATION","TRAIN_ID":"306326","WAITING_SECONDS":"103","WAITING_TIME":"1 min"},{"DESTINATION":"Indian Creek","DIRECTION":"E","EVENT_TIME":"1\/26\/2017 9:58:25 PM","LINE":"BLUE","NEXT_ARR":"10:01:10 PM","STATION":"INDIAN CREEK STATION","TRAIN_ID":"108026","WAITING_SECONDS":"103","WAITING_TIME":"1 min"}]' diff --git a/tests/test_services.py b/tests/test_services.py index a62f5f6..bc6d4b4 100644 --- a/tests/test_services.py +++ b/tests/test_services.py @@ -32,7 +32,7 @@ def test_client_delegates_bus_filters_and_raw_feeds(): assert client.get_bus_vehicle_positions() is client.buses.get_vehicle_positions.return_value assert client.get_bus_trip_updates() is client.buses.get_trip_updates.return_value -def test_bus_service_queries_without_key_and_preserves_bus_contract(monkeypatch): +def test_bus_service_queries_with_key_and_preserves_bus_contract(monkeypatch): feed = gtfs_vehicle_feed( {'route': 1, 'stop_id': 2, 'vehicle_id': 3}, {'route': 4, 'stop_id': 2, 'vehicle_id': 5}, @@ -43,7 +43,7 @@ def test_bus_service_queries_without_key_and_preserves_bus_contract(monkeypatch) ) get = Mock(return_value=response) monkeypatch.setattr(_shared.CACHE, 'get', get) - buses = MartaBusService().get_buses(route=1, stop_id=2, vehicle_id=3) + buses = MartaBusService(api_key='secret').get_buses(route=1, stop_id=2, vehicle_id=3) assert len(buses) == 1 assert isinstance(buses[0], Bus) assert str(buses[0].route) == '1' @@ -63,7 +63,7 @@ def test_bus_service_returns_raw_feed(monkeypatch, method, url): chunks=[feed.SerializeToString()], )) monkeypatch.setattr(_shared.CACHE, 'get', get) - assert getattr(MartaBusService(), method)() == feed + assert getattr(MartaBusService(api_key='secret'), method)() == feed assert get.call_args.args == (url,) @@ -73,24 +73,42 @@ def test_bus_service_retains_unsupported_filter_errors(filters): MartaBusService().get_buses(**filters) -def test_train_service_filters_and_scopes_key_override(monkeypatch, train_response): +def test_train_service_filters_and_uses_configured_key(monkeypatch, train_response): + monkeypatch.setenv('MARTA_API_KEY', 'environment') get = Mock(return_value=FakeResponse(json.loads(train_response))) monkeypatch.setattr(_shared.CACHE, 'get', get) - service = MartaTrainService() - trains = service.get_trains(line='blue', station='Indian Creek Station', direction='east', api_key='override') + service = MartaTrainService(api_key='explicit') + trains = service.get_trains( + line='blue', station='Indian Creek Station', direction='east' + ) assert len(trains) == 1 assert isinstance(trains[0], Train) assert trains[0].direction == 'E' - assert get.call_args.kwargs['params'] == {'apiKey': 'override'} + assert get.call_args.kwargs['params'] == {'apiKey': 'explicit'} service.get_trains() - assert get.call_args.kwargs['params'] == {'apiKey': 'environment'} - MartaTrainService(api_key='explicit').get_trains() assert get.call_args.kwargs['params'] == {'apiKey': 'explicit'} -def test_train_service_requires_key(monkeypatch): +@pytest.mark.parametrize('environment_key', [None, 'environment']) +def test_train_service_requires_explicit_key(monkeypatch, environment_key): + if environment_key is None: + monkeypatch.delenv('MARTA_API_KEY', raising=False) + else: + monkeypatch.setenv('MARTA_API_KEY', environment_key) + get = Mock(side_effect=AssertionError('Missing key must not fetch')) + monkeypatch.setattr(_shared.CACHE, 'get', get) with pytest.raises(APIKeyError): - MartaTrainService().get_trains() + MartaTrainService() + get.assert_not_called() + + +@pytest.mark.parametrize('method', ['get_buses', 'get_vehicle_positions', 'get_trip_updates']) +def test_bus_service_requires_key_before_fetching(monkeypatch, method): + get = Mock(side_effect=AssertionError('Missing key must not fetch')) + monkeypatch.setattr(_shared.CACHE, 'get', get) + with pytest.raises(APIKeyError): + getattr(MartaBusService(), method)() + get.assert_not_called() def test_train_service_retains_invalid_direction_error(monkeypatch):