diff --git a/README.md b/README.md index 048818a..810acd3 100644 --- a/README.md +++ b/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. --- diff --git a/scripts/refresh_data.py b/scripts/refresh_data.py index 41646df..ecba9b4 100644 --- a/scripts/refresh_data.py +++ b/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)") diff --git a/tests/test_refresh_data.py b/tests/test_refresh_data.py new file mode 100644 index 0000000..688776f --- /dev/null +++ b/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