Tracked nr_requests.py and added fetch_nr_timetable_files.

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,
+)

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())

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",
+    )