Merge pull request #10 from PR0M3TH3AN/codex/create-common.py-with-reusable-functions

Refactor helper functions

scripts/bitcoin_m2_chart.py

@@ -19,21 +19,7 @@ import matplotlib.dates as mdates
 import matplotlib.pyplot as plt
 import pandas as pd
 from dateutil.relativedelta import relativedelta
-import sqlite3
+from scripts.common import fetch_series_db, save_figure
-
-
-def save_figure(fig: plt.Figure, output: str | None, script_path: str) -> Path:
-    """Save ``fig`` to ``output`` or a timestamped path under ``outputs/``."""
-    out_path = (
-        Path(output)
-        if output is not None
-        else Path("outputs")
-        / f"{Path(script_path).stem}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.png"
-    )
-    out_path.parent.mkdir(parents=True, exist_ok=True)
-    fig.savefig(out_path)
-    print(f"Saved figure to {out_path}")
-    return out_path
 
 
 def fetch_series(
@@ -44,25 +30,9 @@ def fetch_series(
     db_path: str | Path = Path("data/fred.db"),
 ) -> Tuple[pd.DataFrame, pd.DataFrame]:
     """Load Bitcoin and M2 series from a local SQLite DB."""
-    db_path = Path(db_path)
+    df = fetch_series_db([btc_series, m2_series], start, end, db_path)
-    if not db_path.exists():
+    btc = df[[btc_series]].rename(columns={btc_series: "value"})
-        raise FileNotFoundError(
+    m2 = df[[m2_series]].rename(columns={m2_series: "value"})
-            f"{db_path} not found. Run scripts/refresh_data.py on a machine with internet, commit the DB, then retry."
-        )
-
-    with sqlite3.connect(db_path) as conn:
-        btc = pd.read_sql(
-            f"SELECT date, {btc_series} AS value FROM {btc_series}",
-            conn,
-            parse_dates=["date"],
-            index_col="date",
-        ).loc[start:end]
-        m2 = pd.read_sql(
-            f"SELECT date, {m2_series} AS value FROM {m2_series}",
-            conn,
-            parse_dates=["date"],
-            index_col="date",
-        ).loc[start:end]
     return btc, m2
 
 

scripts/common.py

@@ -0,0 +1,74 @@
+from __future__ import annotations
+
+from datetime import datetime
+from pathlib import Path
+from typing import Iterable
+
+import matplotlib.pyplot as plt
+import pandas as pd
+import sqlite3
+
+
+def save_figure(fig: plt.Figure, output: str | None, script_path: str) -> Path:
+    """Save ``fig`` to ``output`` or ``outputs/`` with timestamp."""
+    out_path = (
+        Path(output)
+        if output is not None
+        else Path("outputs")
+        / f"{Path(script_path).stem}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.png"
+    )
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    fig.savefig(out_path)
+    print(f"Saved figure to {out_path}")
+    return out_path
+
+
+def fetch_series_db(
+    series: Iterable[str],
+    start: datetime,
+    end: datetime,
+    db_path: str | Path = Path("data/fred.db"),
+) -> pd.DataFrame:
+    """Load one or more FRED series from a local SQLite DB."""
+    db_path = Path(db_path)
+    if not db_path.exists():
+        raise FileNotFoundError(
+            f"{db_path} not found. Run scripts/refresh_data.py on a machine with internet, commit the DB, then retry."
+        )
+
+    frames: list[pd.DataFrame] = []
+    with sqlite3.connect(db_path) as conn:
+        for name in series:
+            df = pd.read_sql(
+                f"SELECT date, {name} FROM {name}",
+                conn,
+                parse_dates=["date"],
+                index_col="date",
+            ).loc[start:end]
+            df.rename(columns={name: name}, inplace=True)
+            frames.append(df)
+
+    if not frames:
+        raise ValueError("No series loaded")
+
+    return pd.concat(frames, axis=1).sort_index()
+
+
+def validate_dataframe(df: pd.DataFrame, name: str) -> None:
+    """Ensure dataframe has a ``value`` column with data and DatetimeIndex."""
+    if df.empty:
+        raise ValueError(f"{name} dataframe is empty")
+    if "value" not in df.columns:
+        raise ValueError(f"{name} missing 'value' column")
+    if df["value"].isna().any():
+        raise ValueError(f"{name} contains null values")
+    if not isinstance(df.index, pd.DatetimeIndex):
+        raise ValueError("Indexes must be DatetimeIndex")
+
+
+def validate_overlap(df1: pd.DataFrame, df2: pd.DataFrame) -> None:
+    """Verify that two series share a non-empty overlapping date range."""
+    start = max(df1.index.min(), df2.index.min())
+    end = min(df1.index.max(), df2.index.max())
+    if start >= end:
+        raise ValueError("Series have no overlapping date range")

scripts/custom_chart.py

@@ -17,21 +17,7 @@ from typing import Iterable
 import matplotlib.dates as mdates
 import matplotlib.pyplot as plt
 import pandas as pd
-import sqlite3
+from scripts.common import fetch_series_db, save_figure
-
-
-def save_figure(fig: plt.Figure, output: str | None, script_path: str) -> Path:
-    """Save ``fig`` to ``output`` or a timestamped path under ``outputs/``."""
-    out_path = (
-        Path(output)
-        if output is not None
-        else Path("outputs")
-        / f"{Path(script_path).stem}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.png"
-    )
-    out_path.parent.mkdir(parents=True, exist_ok=True)
-    fig.savefig(out_path)
-    print(f"Saved figure to {out_path}")
-    return out_path
 
 
 def fetch_series_multi(
@@ -41,29 +27,7 @@ def fetch_series_multi(
     db_path: str | Path = Path("data/fred.db"),
 ) -> pd.DataFrame:
     """Load one or more FRED series from a local SQLite DB."""
-    db_path = Path(db_path)
+    return fetch_series_db(series, start, end, db_path)
-    if not db_path.exists():
-        raise FileNotFoundError(
-            f"{db_path} not found. Run scripts/refresh_data.py on a machine with internet, commit the DB, then retry."
-        )
-
-    frames: list[pd.DataFrame] = []
-    with sqlite3.connect(db_path) as conn:
-        for name in series:
-            df = pd.read_sql(
-                f"SELECT date, {name} FROM {name}",
-                conn,
-                parse_dates=["date"],
-                index_col="date",
-            ).loc[start:end]
-            df.rename(columns={name: name}, inplace=True)
-            frames.append(df)
-
-    if not frames:
-        raise ValueError("No series loaded")
-
-    combined = pd.concat(frames, axis=1).sort_index()
-    return combined
 
 
 def plot_series(df: pd.DataFrame) -> plt.Figure:

scripts/lagged_oil_unrate_chart_styled.py

@@ -17,24 +17,14 @@ from pathlib import Path
 import matplotlib.dates as mdates
 import matplotlib.pyplot as plt
 import matplotlib.ticker as mticker
-import numpy as np
 import pandas as pd
 from dateutil.relativedelta import relativedelta
-import sqlite3
+from scripts.common import (
-
+    fetch_series_db,
-
+    save_figure,
-def save_figure(fig: plt.Figure, output: str | None, script_path: str) -> Path:
+    validate_dataframe,
-    """Save ``fig`` to ``output`` or a timestamped path under ``outputs/``."""
+    validate_overlap,
-    out_path = (
+)
-        Path(output)
-        if output is not None
-        else Path("outputs")
-        / f"{Path(script_path).stem}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.png"
-    )
-    out_path.parent.mkdir(parents=True, exist_ok=True)
-    fig.savefig(out_path)
-    print(f"Saved figure to {out_path}")
-    return out_path
 
 
 # ──────────────────────────────────────────────────────────────────────────
@@ -51,26 +41,9 @@ def fetch_series(
     * convert UNRATE to month-end timestamps
     * average WTI daily prices to month-end
     """
-    db_path = Path(db_path)
+    df = fetch_series_db(["UNRATE", "DCOILWTICO"], start, end, db_path)
-    if not db_path.exists():
+    unrate = df[["UNRATE"]].rename(columns={"UNRATE": "value"})
-        raise FileNotFoundError(
+    oil = df[["DCOILWTICO"]].rename(columns={"DCOILWTICO": "value"})
-            f"{db_path} not found.  Run scripts/refresh_data.py on a "
-            "machine with internet, commit the new DB, then retry."
-        )
-
-    with sqlite3.connect(db_path) as conn:
-        unrate = pd.read_sql(
-            "SELECT date, UNRATE AS value FROM UNRATE",
-            conn,
-            parse_dates=["date"],
-            index_col="date",
-        ).loc[start:end]
-        oil = pd.read_sql(
-            "SELECT date, DCOILWTICO AS value FROM DCOILWTICO",
-            conn,
-            parse_dates=["date"],
-            index_col="date",
-        ).loc[start:end]
 
     # Resample WTI to month-end average
     oil_monthly = oil.resample("M").mean()
@@ -80,34 +53,10 @@ def fetch_series(
 
 
 def validate_series(unrate: pd.DataFrame, oil: pd.DataFrame) -> None:
-    """
+    """Run basic sanity checks on the two series."""
-    Sanity-check the fetched series before plotting.
+    validate_dataframe(unrate, "UNRATE")
-
+    validate_dataframe(oil, "DCOILWTICO")
-    Both dataframes must contain a ``value`` column with no missing values
+    validate_overlap(unrate, oil)
-    and share a non-empty overlapping date range.
-    """
-    if unrate.empty:
-        raise ValueError("UNRATE dataframe is empty")
-    if oil.empty:
-        raise ValueError("DCOILWTICO dataframe is empty")
-
-    for name, df in [("UNRATE", unrate), ("DCOILWTICO", oil)]:
-        if "value" not in df.columns:
-            raise ValueError(f"{name} missing 'value' column")
-        if df["value"].isna().any():
-            raise ValueError(f"{name} contains null values")
-
-    if not isinstance(unrate.index, pd.DatetimeIndex) or not isinstance(
-        oil.index, pd.DatetimeIndex
-    ):
-        raise ValueError("Indexes must be DatetimeIndex")
-
-    # Replace strict range check with overlap check
-    overlap_start = max(unrate.index.min(), oil.index.min())
-    overlap_end = min(unrate.index.max(), oil.index.max())
-
-    if overlap_start >= overlap_end:
-        raise ValueError("Series have no overlapping date range")
 
 
 # ──────────────────────────────────────────────────────────────────────────