Initial commit

agenda/spacexdata.py

@@ -0,0 +1,139 @@
+import json
+import os
+import sys
+from datetime import datetime, timezone
+from typing import Any
+
+import dateutil
+import dateutil.parser
+import requests
+
+data_dir = "/home/edward/lib/data"
+spacex_dir = os.path.join(data_dir, "spacex")
+
+now = datetime.now()
+here = dateutil.tz.tzlocal()
+do_refresh = len(sys.argv) > 1 and sys.argv[1] == "--refresh"
+
+Launch = dict[str, str | datetime | list[Any]]
+
+
+def filename_timestamp(filename: str) -> tuple[datetime, str] | None:
+    """Get datetime from filename."""
+    try:
+        ts = datetime.strptime(filename, "%Y-%m-%d_%H:%M:%S.json")
+    except ValueError:
+        return None
+    return (ts, filename)
+
+
+def next_spacex_launch_api(limit: int) -> list[Launch]:
+    """Get the next upcoming launches from the API."""
+    filename = os.path.join(spacex_dir, now.strftime("%Y-%m-%d_%H:%M:%S.json"))
+    url = "https://api.spacexdata.com/v4/launches/upcoming"
+
+    params: dict[str, str | int] = dict(tbd="false", limit=limit)
+    r = requests.get(url, params=params)
+    open(filename, "w").write(r.text)
+    launch_list = r.json()
+    if "error" in launch_list:
+        print(r.url)
+        print(launch_list)
+        raise ValueError
+    return [
+        parse_get_next_spacex_launch(launch)
+        for launch in launch_list
+        if launch["date_precision"] not in ("month", "quarter", "half")
+    ]
+
+
+def get_spacex_payloads() -> dict[str, str]:
+    """Load SpaceX payloads or refresh if missing."""
+    filename = os.path.join(spacex_dir, "payloads.json")
+    if not os.path.exists(filename):
+        refresh_payloads()
+    data = json.load(open(filename))
+    # print(json.dumps(data, indent=2))
+    for p in data:
+        if "type" in p:
+            continue
+        print(json.dumps(p, indent=2))
+    return {p["id"]: p["type"] for p in data}
+
+
+def get_spacex_rockets() -> dict[str, str]:
+    """Load mapping of rocket ID to rocket name."""
+    filename = os.path.join(spacex_dir, "rockets.json")
+    return {r["id"]: r["name"] for r in json.load(open(filename))}
+
+
+def refresh_payloads() -> None:
+    """Download list of payloads and save."""
+    url = "https://api.spacexdata.com/v4/payloads"
+    filename = os.path.join(spacex_dir, "payloads.json")
+    r = requests.get(url)
+    open(filename, "w").write(r.text)
+
+
+rocket_map = get_spacex_rockets()
+payloads = get_spacex_payloads()
+
+
+def parse_get_next_spacex_launch(launch: dict[str, Any]) -> Launch:
+    global payloads
+    date_utc = dateutil.parser.parse(launch["date_utc"])
+    date_utc.replace(tzinfo=timezone.utc)
+
+    # TODO: payload
+
+    need_refresh = any(p not in payloads for p in launch["payloads"])
+    if need_refresh:
+        refresh_payloads()
+        payloads = get_spacex_payloads()
+        assert all(p in payloads for p in launch["payloads"])
+
+    return {
+        "rocket": rocket_map[launch["rocket"]],
+        "name": launch["name"],
+        "when": date_utc.astimezone(here),
+        "date_precision": launch["date_precision"],
+        "payloads": [payloads.get(p, "[unknown payload]") for p in launch["payloads"]],
+    }
+
+
+def get_most_recent_existing(existing):
+    for _, f in existing:
+        filename = os.path.join(spacex_dir, f)
+        launch_list = json.load(open(filename))
+        if "error" in launch_list:
+            print(launch_list)
+            continue
+
+        return launch_list
+
+
+def get_next_spacex_launch(limit: int) -> list[Launch]:
+    filename = os.path.join(data_dir, "next_spacex_launch.json")
+    spacex_dir = os.path.join(data_dir, "spacex")
+
+    existing = [x for x in (filename_timestamp(f) for f in os.listdir(spacex_dir)) if x]
+
+    existing.sort(reverse=True)
+
+    if do_refresh or not existing or (now - existing[0][0]).seconds > 3600:  # one hour
+        try:
+            return next_spacex_launch_api(limit=limit)
+        except ValueError:
+            print("*** SpaceX next launch error ***")
+
+    for _, f in existing:
+        filename = os.path.join(spacex_dir, f)
+        launch_list = json.load(open(filename))
+        if "error" in launch_list:
+            print(launch_list)
+            continue
+        return [
+            parse_get_next_spacex_launch(launch)
+            for launch in launch_list
+            if launch["date_precision"] not in ("month", "quarter", "half")
+        ]