Merge pull request #15 from PR0M3TH3AN/codex/extend-argument-parsing-for-plotting-scripts

Add figure size and display options

scripts/bitcoin_m2_chart.py

@@ -59,6 +59,10 @@ def plot_bitcoin_m2(
     m2: pd.DataFrame,
     offset_days: int,
     extend_years: int,
+    *,
+    width: float = 12,
+    height: float = 6,
+    dpi: int | None = None,
 ) -> plt.Figure:
     """Plot Bitcoin price and shifted M2 on dual axes."""
     # Determine overlapping range
@@ -70,7 +74,7 @@ def plot_bitcoin_m2(
 
     m2_shifted = m2_common.shift(offset_days)
 
-    fig, ax1 = plt.subplots(figsize=(12, 6))
+    fig, ax1 = plt.subplots(figsize=(width, height), dpi=dpi)
     ax1.plot(
         btc_common.index, btc_common["value"], label="Bitcoin (LHS)", color="#1f77b4"
     )
@@ -137,6 +141,10 @@ def main(argv: list[str] | None = None) -> plt.Figure:
     p.add_argument(
         "--output", type=str, default=None, help="optional path to save the figure"
     )
+    p.add_argument("--width", type=float, default=12.0, help="figure width in inches")
+    p.add_argument("--height", type=float, default=6.0, help="figure height in inches")
+    p.add_argument("--dpi", type=int, default=None, help="figure DPI")
+    p.add_argument("--no-show", action="store_true", help="do not display the figure")
     args = p.parse_args(argv)
 
     start_dt = datetime.fromisoformat(args.start)
@@ -144,10 +152,18 @@ def main(argv: list[str] | None = None) -> plt.Figure:
 
     btc, m2 = fetch_series(start_dt, end_dt, args.btc_series, args.m2_series, args.db)
     validate_series(btc, m2, args.btc_series, args.m2_series)
-    fig = plot_bitcoin_m2(btc, m2, args.offset, args.extend_years)
+    fig = plot_bitcoin_m2(
+        btc,
+        m2,
+        args.offset,
+        args.extend_years,
+        width=args.width,
+        height=args.height,
+        dpi=args.dpi,
+    )
 
     save_figure(fig, args.output, __file__)
-    if argv is None:
+    if argv is None and not args.no_show:
         fig.show()
     return fig
 

scripts/lagged_oil_unrate_chart_styled.py

@@ -82,6 +82,10 @@ def plot_lagged(
     _start_date: datetime,
     _end_date: datetime,
     extend_years: int,
+    *,
+    width: float = 14,
+    height: float = 7,
+    dpi: int | None = None,
 ) -> plt.Figure:
     """
     Render the dual-axis chart with fixed scales and log ticks on oil.
@@ -107,7 +111,7 @@ def plot_lagged(
     # Apply lag/lead
     oil_shifted = oil_common.shift(offset_months)
 
-    fig, ax1 = plt.subplots(figsize=(14, 7))
+    fig, ax1 = plt.subplots(figsize=(width, height), dpi=dpi)
     plt.title("Unemployment Rate and US Oil Price", fontsize=20, weight="bold")
     plt.suptitle(
         "Civilian Unemployment Rate and WTI Crude Oil Price",
@@ -193,7 +197,7 @@ def plot_lagged(
 # ──────────────────────────────────────────────────────────────────────────
 # CLI
 # ──────────────────────────────────────────────────────────────────────────
-def main() -> plt.Figure:
+def main(argv: list[str] | None = None) -> plt.Figure:
     p = argparse.ArgumentParser()
     p.add_argument(
         "-o",
@@ -232,7 +236,11 @@ def main() -> plt.Figure:
         default=None,
         help="optional path to save the figure (PNG or PDF)",
     )
-    args = p.parse_args()
+    p.add_argument("--width", type=float, default=14.0, help="figure width in inches")
+    p.add_argument("--height", type=float, default=7.0, help="figure height in inches")
+    p.add_argument("--dpi", type=int, default=None, help="figure DPI")
+    p.add_argument("--no-show", action="store_true", help="do not display the figure")
+    args = p.parse_args(argv)
 
     start_dt = datetime.fromisoformat(args.start)
     end_dt = datetime.fromisoformat(args.end)
@@ -256,9 +264,20 @@ def main() -> plt.Figure:
     oil = oil.loc[oil.index.isin(unrate.index)]
 
     validate_series(unrate, oil)
-    fig = plot_lagged(unrate, oil, args.offset, start_dt, end_dt, args.extend_years)
+    fig = plot_lagged(
+        unrate,
+        oil,
+        args.offset,
+        start_dt,
+        end_dt,
+        args.extend_years,
+        width=args.width,
+        height=args.height,
+        dpi=args.dpi,
+    )
     save_figure(fig, args.output, __file__)
-    fig.show()
+    if argv is None and not args.no_show:
+        fig.show()
     return fig
 
 

tests/test_bitcoin_m2.py

@@ -32,7 +32,19 @@ def test_btc_m2_main(monkeypatch):
         return btc, m2
 
     monkeypatch.setattr("scripts.bitcoin_m2_chart.fetch_series", fake_fetch)
-    fig = btc_m2_main(["--btc-series", "BTC", "--m2-series", "M2"])
+    fig = btc_m2_main([
+        "--btc-series",
+        "BTC",
+        "--m2-series",
+        "M2",
+        "--width",
+        "8",
+        "--height",
+        "4",
+        "--dpi",
+        "150",
+        "--no-show",
+    ])
     assert len(fig.axes) == 2
 
 
@@ -52,6 +64,13 @@ def test_btc_m2_main_saves_output(tmp_path, monkeypatch):
             "M2",
             "--output",
             str(out_file),
+            "--width",
+            "8",
+            "--height",
+            "4",
+            "--dpi",
+            "150",
+            "--no-show",
         ]
     )
     assert out_file.exists()

tests/test_plot.py

@@ -8,7 +8,7 @@ import pandas as pd
 
 sys.path.append(os.path.abspath("."))
 
-from scripts.lagged_oil_unrate_chart_styled import plot_lagged
+from scripts.lagged_oil_unrate_chart_styled import plot_lagged, main as lagged_main
 from scripts.custom_chart import plot_series, main as custom_main
 
 
@@ -64,3 +64,18 @@ def test_custom_main_saves_output(tmp_path):
         ]
     )
     assert out_file.exists()
+
+
+def test_lagged_main_options(monkeypatch):
+    dates = pd.date_range("2021-01-01", periods=3, freq="M")
+    unrate = pd.DataFrame({"value": [1, 2, 3]}, index=dates)
+    oil = pd.DataFrame({"value": [4, 5, 6]}, index=dates)
+
+    monkeypatch.setattr(
+        "scripts.lagged_oil_unrate_chart_styled.fetch_series", lambda *_: (unrate, oil)
+    )
+
+    fig = lagged_main(
+        ["--offset", "1", "--width", "8", "--height", "4", "--dpi", "150", "--no-show"]
+    )
+    assert len(fig.axes) == 2