From fe1dc2b52d72c1084e87ad14b4a410f266eabc24 Mon Sep 17 00:00:00 2001
From: Timothy Willard <9395586+TimothyWillard@users.noreply.github.com>
Date: Wed, 11 Dec 2024 08:58:07 -0500
Subject: [PATCH] Add `--indent` option to `flepimop patch`

For configuration of the outputted yaml.
---
 flepimop/gempyor_pkg/src/gempyor/cli.py       | 23 ++++++++++++++----
 .../tests/cli/test_flepimop_patch_cli.py      | 24 +++++++++++++++++++
 2 files changed, 42 insertions(+), 5 deletions(-)

diff --git a/flepimop/gempyor_pkg/src/gempyor/cli.py b/flepimop/gempyor_pkg/src/gempyor/cli.py
index 70da194d2..45433f4d5 100644
--- a/flepimop/gempyor_pkg/src/gempyor/cli.py
+++ b/flepimop/gempyor_pkg/src/gempyor/cli.py
@@ -1,4 +1,5 @@
-from click import pass_context, Context
+import click
+import yaml
 
 from .shared_cli import (
     config_files_argument,
@@ -20,12 +21,24 @@
 
 
 # add some basic commands to the CLI
-@cli.command(params=[config_files_argument] + list(config_file_options.values()))
-@pass_context
-def patch(ctx: Context = mock_context, **kwargs) -> None:
+@cli.command(
+    params=[config_files_argument]
+    + list(config_file_options.values())
+    + [
+        click.Option(
+            ["--indent"],
+            type=click.IntRange(min=1),
+            required=False,
+            default=2,
+            help="Indentation level for the output YAML.",
+        )
+    ],
+)
+@click.pass_context
+def patch(ctx: click.Context = mock_context, **kwargs) -> None:
     """Merge configuration files"""
     parse_config_files(config, ctx, **kwargs)
-    print(config.dump())
+    print(yaml.dump(yaml.safe_load(config.dump()), indent=kwargs["indent"]))
 
 
 if __name__ == "__main__":
diff --git a/flepimop/gempyor_pkg/tests/cli/test_flepimop_patch_cli.py b/flepimop/gempyor_pkg/tests/cli/test_flepimop_patch_cli.py
index 72d0134b9..800fac3b1 100644
--- a/flepimop/gempyor_pkg/tests/cli/test_flepimop_patch_cli.py
+++ b/flepimop/gempyor_pkg/tests/cli/test_flepimop_patch_cli.py
@@ -74,3 +74,27 @@ def test_patch_seir_parameters_behavior(
     patched_config = yaml.safe_load(result.output)
     assert "seir" in patched_config
     assert patched_config["seir"]["parameters"] == expected_parameters
+
+
+@pytest.mark.parametrize("indent", (2, 4, 6))
+def test_user_provided_indent_size(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch, indent: int
+) -> None:
+    # Setup the test
+    monkeypatch.chdir(tmp_path)
+    config = tmp_path / "config.yml"
+    config.write_text(yaml.dump({"seir": {"parameters": {"beta": {"value": 1.2}}}}))
+
+    # Invoke the command
+    runner = CliRunner()
+    result = runner.invoke(patch, [config.name, "--indent", str(indent)])
+    assert result.exit_code == 0
+
+    # Check the output indentation, manually since `yaml.load` abstracts spacing away
+    for line in result.output.split():
+        stripped_line = line.lstrip()
+        if stripped_line and not stripped_line.startswith("#"):
+            indent_size = len(line) - len(stripped_line)
+            if indent_size > 0:
+                assert indent_size == indent
+                break