Merge pull request #4 from PR0M3TH3AN/codex/add---output-argument-to-chart-scripts

Add optional output save path

.gitignore

@@ -8,3 +8,4 @@ wheelhouse/*.md
 # keep wheel files themselves!
 data/*.bak
 repo-context.txt
+outputs/

README.md

@@ -79,6 +79,7 @@ requested period before rendering the chart.
 | `--start`        | Start date (YYYY-MM-DD)                         |
 | `--end`          | End date (YYYY-MM-DD)                           |
 | `--extend-years` | Years to extend the x-axis beyond the end date  |
+| `--output`       | Save the figure to this path before showing it  |
 
 ## Custom Charts
 

scripts/custom_chart.py

@@ -83,6 +83,12 @@ def main(argv: list[str] | None = None) -> plt.Figure:
         help="end date YYYY-MM-DD",
     )
     p.add_argument("--db", type=str, default="data/fred.db", help="path to SQLite DB")
+    p.add_argument(
+        "--output",
+        type=str,
+        default=None,
+        help="optional path to save the figure (PNG or PDF)",
+    )
     args = p.parse_args(argv)
 
     start_dt = datetime.fromisoformat(args.start)
@@ -91,6 +97,8 @@ def main(argv: list[str] | None = None) -> plt.Figure:
     data = fetch_series_multi(args.series, start_dt, end_dt, args.db)
     fig = plot_series(data)
 
+    if args.output:
+        fig.savefig(args.output)
     if argv is None:
         fig.show()
     return fig

scripts/lagged_oil_unrate_chart_styled.py

@@ -250,6 +250,12 @@ def main() -> None:
         default="data/fred.db",
         help="path to local FRED SQLite file",
     )
+    p.add_argument(
+        "--output",
+        type=str,
+        default=None,
+        help="optional path to save the figure (PNG or PDF)",
+    )
     args = p.parse_args()
 
     start_dt = datetime.fromisoformat(args.start)
@@ -258,6 +264,8 @@ def main() -> None:
     unrate, oil = fetch_series(start_dt, end_dt, args.db)
     validate_series(unrate, oil)
     fig = plot_lagged(unrate, oil, args.offset, start_dt, end_dt, args.extend_years)
+    if args.output:
+        fig.savefig(args.output)
     fig.show()
 
 

tests/test_plot.py

@@ -45,3 +45,22 @@ def test_plot_series_and_main():
     )
     assert len(fig2.axes) == 1
     assert any(fig2.axes[0].get_lines())
+
+
+def test_custom_main_saves_output(tmp_path):
+    out_file = tmp_path / "chart.png"
+    custom_main(
+        [
+            "--series",
+            "UNRATE",
+            "--start",
+            "2020-01-01",
+            "--end",
+            "2020-03-01",
+            "--db",
+            "data/fred.db",
+            "--output",
+            str(out_file),
+        ]
+    )
+    assert out_file.exists()