From fb273198a0d3905f6c0568c39b508127f6ed6d58 Mon Sep 17 00:00:00 2001
From: Gregor Sturm <gregor.sturm@boehringer-ingelheim.com>
Date: Fri, 29 Nov 2024 14:19:10 +0100
Subject: [PATCH] Close #65

---
 CHANGELOG.md             |  1 +
 src/dso/get_config.py    |  4 ++++
 tests/test_get_config.py | 38 ++++++++++++++++++++++++++++++++++++++
 3 files changed, 43 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 3c465a0..9fb6389 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -16,6 +16,7 @@ and this project adheres to [Semantic Versioning][].
 -   Make it possible to override watermarks/disclaimers with a simple `null` ([#69](https://github.com/Boehringer-Ingelheim/dso/pull/69)).
 -   Compile *all* configs on `dso repro`, not just the ones relvant to the specified stage. This is required because we don't
     know which other stages dvc might compile ([#69](https://github.com/Boehringer-Ingelheim/dso/pull/69)).
+-   Make `get-config` compatible with dvc matrix stages ([#69](https://github.com/Boehringer-Ingelheim/dso/pull/69)).
 
 ### Template updates
 
diff --git a/src/dso/get_config.py b/src/dso/get_config.py
index a4074f8..2cfa374 100644
--- a/src/dso/get_config.py
+++ b/src/dso/get_config.py
@@ -99,6 +99,7 @@ def get_config(stage: str, *, all: bool = False, skip_compile: bool = False) ->
 
         # We want to include parameters mentioned in either `params`, `deps`, `outs`.
         # The parameters in `deps`/`outs` are encapsulated in `${ <param> }`
+        is_matrix_stage = "matrix" in dvc_stage_config
         keep_params = set(dvc_stage_config.get("params", []))
         dvc_param_pat = re.compile(r"\$\{\s*(.*?)\s*\}")
         for dep in dvc_stage_config.get("deps", []):
@@ -112,6 +113,9 @@ def get_config(stage: str, *, all: bool = False, skip_compile: bool = False) ->
             f"Only including the following parameters which are listed in `dvc.yaml`: [green]{', '.join(keep_params)}"
         )
 
+        if is_matrix_stage:
+            keep_params = {p for p in keep_params if not (p.startswith("item.") or p == "item")}
+
         return _filter_nested_dict(config, keep_params)
 
 
diff --git a/tests/test_get_config.py b/tests/test_get_config.py
index 1671c7e..e607ce3 100644
--- a/tests/test_get_config.py
+++ b/tests/test_get_config.py
@@ -142,6 +142,44 @@ def test_get_config_order(dso_project):
     assert list(config["B1"]) == ["C", "A", "D", "B", "Z", "X"]
 
 
+def test_get_config_matrix(dso_project):
+    """Test that get-config is compatible with dvc's matrix feature"""
+    chdir(dso_project)
+    stage = dso_project / "mystage"
+    stage.mkdir()
+    (stage / "params.in.yaml").touch()
+
+    (dso_project / "params.in.yaml").write_text(
+        dedent(
+            """\
+            matrix_param: ['p1', 'p2']
+            A: "aaa"
+            B: "bbb"
+            """
+        )
+    )
+
+    (stage / "dvc.yaml").write_text(
+        dedent(
+            """\
+            stages:
+               mystage01:
+                 matrix:
+                    mp: ${ matrix_param }
+                 params:
+                   - A
+                   - item.mp
+                 outs:
+                   - output/${ item.mp }
+                 cmd: "echo Hello World!"
+            """
+        )
+    )
+
+    config = get_config("mystage")
+    assert list(config) == ["A"]
+
+
 def test_get_config_path_relative_to_root_dir(quarto_stage):
     chdir(quarto_stage)
     config1 = get_config("quarto_stage")