diff --git a/src/national_rail_timetable/__init__.py b/src/national_rail_timetable/__init__.py index e69de29..79131c9 100644 --- a/src/national_rail_timetable/__init__.py +++ b/src/national_rail_timetable/__init__.py @@ -0,0 +1,12 @@ +""" +Expose classes and functions to avoid unneeded nesting. +""" +# pyright: reportUnusedImport=false +# ruff: noqa F401 + +# Imports +from national_rail_timetable.nr_requests import ( + NRConfig, + fetch_nr_token, + fetch_nr_timetable_files, +) diff --git a/src/national_rail_timetable/__main__.py b/src/national_rail_timetable/__main__.py index 5e4bdc5..765508a 100644 --- a/src/national_rail_timetable/__main__.py +++ b/src/national_rail_timetable/__main__.py @@ -1,4 +1,4 @@ -from national_rail_timetable.nr_requests import fetch_nr_token +from national_rail_timetable.nr_requests import fetch_nr_token, fetch_nr_timetable_files print(fetch_nr_token()) - +print(fetch_nr_timetable_files()) diff --git a/src/national_rail_timetable/nr_requests.py b/src/national_rail_timetable/nr_requests.py new file mode 100644 index 0000000..6756829 --- /dev/null +++ b/src/national_rail_timetable/nr_requests.py @@ -0,0 +1,100 @@ +""" +Functions for authentication / timetable requests with National Rail. +""" + +# Imports +from dataclasses import dataclass +import json +import os +from zipfile import ZipFile +from io import BytesIO +from time import sleep +import requests +from requests import Response + + +# Init. +@dataclass +class NRConfig: + username: str + password: str + + @classmethod + def from_env(cls) -> "NRConfig": + try: + return cls( + username=os.environ["NR_USER"], + password=os.environ["NR_PASS"], + ) + except KeyError: + raise KeyError( + "Not all necessary environment variables are available. " + + "Make sure NR_USER and NR_PASS are defined at config generation time. " + + "Alternatively, provide an NRConfig instance defined manually. " + ) + + +# Functions +def fetch_nr_token( + config: NRConfig | None = None, # pyright: ignore[reportRedeclaration] + attempts: int = 3, +) -> str: + config: NRConfig = config if config is not None else NRConfig.from_env() + response: Response = requests.post( + url="https://opendata.nationalrail.co.uk/authenticate", + headers={"Content-Type": "application/x-www-form-urlencoded"}, + data=config.__dict__, + ) + if not response.ok: + if attempts > 1: + sleep(3) + return fetch_nr_token( + config=config, + attempts=attempts - 1, + ) + raise requests.HTTPError( + "Failed to fetch token using credentials provided in config. Response: " + + f"[{response.status_code}] " + + response.content.decode() + ) + content: str = response.content.decode() + assert isinstance(token := json.loads(s=content)["token"], str) # pyright: ignore[reportAny] + return token + + +def fetch_nr_timetable_files( + config: NRConfig | None = None, # pyright: ignore[reportRedeclaration] + token: str | None = None, # pyright: ignore[reportRedeclaration] + attempts: int = 1, +) -> ZipFile: + config: NRConfig = config if config is not None else NRConfig.from_env() + token: str = ( + token + if token is not None + else fetch_nr_token( + config=config, + attempts=attempts, + ) + ) + response: Response = requests.get( + url="https://opendata.nationalrail.co.uk/api/staticfeeds/3.0/timetable", + headers={"X-Auth-Token": token, "Content-Type": "application/json"}, + ) + if not response.ok: + if attempts > 1: + sleep(3) + return fetch_nr_timetable_files( + config=config, + token=token, + attempts=attempts - 1, + ) + raise requests.HTTPError( + "Failed to fetch timetable files using token and config provided. Response: " + + f"[{response.status_code}] " + + response.content.decode() + ) + zipped_bytes: bytes = response.content + return ZipFile( + BytesIO(initial_bytes=zipped_bytes), + mode="r", + )