From 06c84fb20ec3c9237741b835db210b2faf075b37 Mon Sep 17 00:00:00 2001
From: Tom <thomas.hodson@ecmwf.int>
Date: Mon, 24 Mar 2025 15:28:06 +0000
Subject: [PATCH] Fix selection bug to require that arguments be consumed by a
 branch

---
 src/python/qubed/Qube.py |  1 +
 tests/test_selection.py  | 27 +++++++++++++++++++++++++++
 2 files changed, 28 insertions(+)
 create mode 100644 tests/test_selection.py

diff --git a/src/python/qubed/Qube.py b/src/python/qubed/Qube.py
index 9217a02..90352e3 100644
--- a/src/python/qubed/Qube.py
+++ b/src/python/qubed/Qube.py
@@ -261,6 +261,7 @@ class Qube:
         selection: dict[str, str | list[str]],
         mode: Literal["strict", "relaxed"] = "relaxed",
         prune=True,
+        consume=True,
     ) -> "Qube":
         # make all values lists
         selection = {k: v if isinstance(v, list) else [v] for k, v in selection.items()}
diff --git a/tests/test_selection.py b/tests/test_selection.py
new file mode 100644
index 0000000..019b848
--- /dev/null
+++ b/tests/test_selection.py
@@ -0,0 +1,27 @@
+from qubed import Qube
+
+q = Qube.from_dict(
+    {
+        "class=od": {
+            "expver=0001": {"param=1": {}, "param=2": {}},
+            "expver=0002": {"param=1": {}, "param=2": {}},
+        },
+        "class=rd": {"param=1": {}, "param=2": {}, "param=3": {}},
+    }
+)
+
+
+def test_consumption():
+    assert q.select({"expver": "0001"}) == Qube.from_dict(
+        {"class=od": {"expver=0001": {"param=1": {}, "param=2": {}}}}
+    )
+
+
+def test_consumption_off():
+    expected = Qube.from_dict(
+        {
+            "class=od": {"expver=0001": {"param=1": {}, "param=2": {}}},
+            "class=rd": {"param=1": {}, "param=2": {}, "param=3": {}},
+        }
+    )
+    assert q.select({"expver": "0001"}, consume=False) == expected