Merge pull request #6 from PR0M3TH3AN/codex/modify-refresh_data.py-arguments-and-tests

Enable custom series download

README.md

@@ -55,7 +55,7 @@ python_charts/
 ### Refresh FRED data
 
 ```bash
-./startup.sh scripts/refresh_data.py
+./startup.sh python scripts/refresh_data.py --series UNRATE CPIAUCSL --start 1960-01-01
 ```
 
 Downloads UNRATE and WTI series (from 1948 to today) into `data/fred.db`.
@@ -128,7 +128,7 @@ python-dateutil>=2.8.1
 ## 🛡️ Troubleshooting
 
 * **ModuleNotFoundError: distutils**: resolved by `pip install --upgrade setuptools`.
-* **Database not found**: run `./startup.sh scripts/refresh_data.py`.
+* **Database not found**: run `./startup.sh python scripts/refresh_data.py`.
 * **Permission errors**: ensure write access to the project directory.
 
 ---

scripts/refresh_data.py

@@ -1,33 +1,51 @@
 #!/usr/bin/env python3
-"""
-refresh_data.py
----------------
-Download UNRATE and DCOILWTICO from FRED and store them in SQLite
-(data/fred.db).  Run on a connected machine; commit the DB afterwards.
+"""refresh_data.py
+-----------------
+Download one or more FRED series into ``data/fred.db``.
+
+Run this script on a machine with internet access and commit the updated
+database afterwards.
+
+Example:
+    python scripts/refresh_data.py --series UNRATE CPIAUCSL --start 1960-01-01
 """
 
+from __future__ import annotations
+
 from datetime import date
+import argparse
 import pathlib
 import sqlite3
 
-import pandas as pd
 from pandas_datareader.data import DataReader
+import pandas as pd
 
 HERE = pathlib.Path(__file__).resolve().parent
 DATA_PATH = HERE.parent / "data"
 DB_FILE = DATA_PATH / "fred.db"
 
-START = "1948-01-01"          # earliest UNRATE observation
-END = date.today().isoformat()
-SERIES = ("UNRATE", "DCOILWTICO")
+DEFAULT_START = "1948-01-01"  # earliest UNRATE observation
+DEFAULT_END = date.today().isoformat()
+DEFAULT_SERIES = ("UNRATE", "DCOILWTICO")
 
 
-def main() -> None:
+def main(argv: list[str] | None = None) -> None:
+    p = argparse.ArgumentParser(description="Download FRED series to SQLite")
+    p.add_argument(
+        "--series",
+        nargs="+",
+        default=list(DEFAULT_SERIES),
+        help="FRED series names",
+    )
+    p.add_argument("--start", type=str, default=DEFAULT_START, help="start date")
+    p.add_argument("--end", type=str, default=DEFAULT_END, help="end date")
+    args = p.parse_args(argv)
+
     DATA_PATH.mkdir(exist_ok=True)
     with sqlite3.connect(DB_FILE) as conn:
-        for s in SERIES:
+        for s in args.series:
             print(f"↯ downloading {s} …")
-            df = DataReader(s, "fred", START, END)
+            df = DataReader(s, "fred", args.start, args.end)
             df.to_sql(s, conn, if_exists="replace", index_label="date")
             print(f"✓ wrote {s}: {len(df):,} rows")
     print(f"\nDone.  SQLite DB at {DB_FILE} ({DB_FILE.stat().st_size/1024:.0f} KB)")

tests/test_refresh_data.py

@@ -0,0 +1,44 @@
+import os
+import sys
+import sqlite3
+from pathlib import Path
+
+import pandas as pd
+
+sys.path.append(os.path.abspath("."))
+
+import scripts.refresh_data as refresh_data
+
+
+def test_refresh_data_writes_series(tmp_path, monkeypatch):
+    calls = []
+
+    def fake_reader(name, src, start, end):
+        calls.append(name)
+        dates = pd.date_range(start, periods=1)
+        return pd.DataFrame({name: [1]}, index=dates)
+
+    monkeypatch.setattr(refresh_data, "DataReader", fake_reader)
+    monkeypatch.setattr(refresh_data, "DATA_PATH", tmp_path)
+    db_file = tmp_path / "fred.db"
+    monkeypatch.setattr(refresh_data, "DB_FILE", db_file)
+
+    refresh_data.main(
+        [
+            "--series",
+            "A",
+            "B",
+            "--start",
+            "2000-01-01",
+            "--end",
+            "2000-01-02",
+        ]
+    )
+
+    assert set(calls) == {"A", "B"}
+    with sqlite3.connect(db_file) as conn:
+        tables = {
+            row[0]
+            for row in conn.execute("SELECT name FROM sqlite_master WHERE type='table'")
+        }
+    assert {"A", "B"} <= tables