diff --git a/.gitignore b/.gitignore index 57dd0ef..1408f19 100644 --- a/.gitignore +++ b/.gitignore @@ -8,3 +8,4 @@ wheelhouse/*.md # keep wheel files themselves! data/*.bak repo-context.txt +outputs/ diff --git a/README.md b/README.md index aa803a7..048818a 100644 --- a/README.md +++ b/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 diff --git a/scripts/custom_chart.py b/scripts/custom_chart.py index 7bee02b..9390a9f 100644 --- a/scripts/custom_chart.py +++ b/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 diff --git a/scripts/lagged_oil_unrate_chart_styled.py b/scripts/lagged_oil_unrate_chart_styled.py index 95978e0..d846d41 100644 --- a/scripts/lagged_oil_unrate_chart_styled.py +++ b/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() diff --git a/tests/test_plot.py b/tests/test_plot.py index 791b2b2..e5246e3 100644 --- a/tests/test_plot.py +++ b/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()