Stop support for list of "take" protocol (#1517)

* Stop support for list of "take" protocol * Try to bypass the check of URL * Debug * Debug
XENONnT · Jan 10, 2025 · f993f0b · f993f0b
1 parent 69265ff
commit f993f0b
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 11 deletions.
diff --git a/straxen/config/protocols.py b/straxen/config/protocols.py
@@ -65,6 +65,8 @@ def read_json(content: str, **kwargs):
 @URLConfig.register("take")
 def get_key(container: Container, take=None, **kwargs):
     """Return a single element of a container."""
+    if not isinstance(container, dict):
+        raise ValueError(f"Container is not a dict but a {type(container)}")
     if take is None:
         return container
     if not isinstance(take, list):

diff --git a/tests/test_url_config.py b/tests/test_url_config.py
@@ -25,6 +25,12 @@ def generate_random(_):
     return random.random()
 
 
+@straxen.URLConfig.register("range")
+def generate_range(length):
+    length = int(length)
+    return np.arange(length)
+
+
 @straxen.URLConfig.register("unpicklable")
 def return_lamba(_):
     return lambda x: x
@@ -118,9 +124,9 @@ def test_cmt_protocol(self):
         self.assertTrue(abs(p.test_config - 219203.49884000001) < 1e-2)
 
     def test_json_protocol(self):
-        self.st.set_config({"test_config": "json://[1,2,3]"})
+        self.st.set_config({"test_config": 'json://{"a":0}'})
         p = self.st.get_single_plugin(nt_test_run_id, "test_data")
-        self.assertEqual(p.test_config, [1, 2, 3])
+        self.assertEqual(p.test_config, {"a": 0})
 
     def test_format_protocol(self):
         self.st.set_config({"test_config": "format://{run_id}?run_id=plugin.run_id"})
@@ -136,15 +142,16 @@ def test_fsspec_protocol(self):
         p = self.st.get_single_plugin(nt_test_run_id, "test_data")
         self.assertEqual(p.test_config, 999)
 
-    def test_chained(self):
-        self.st.set_config({"test_config": "take://json://[1,2,3]?take=0"})
-        p = self.st.get_single_plugin(nt_test_run_id, "test_data")
-        self.assertEqual(p.test_config, 1)
-
     def test_take_nested(self):
-        self.st.set_config({"test_config": 'take://json://{"a":[1,2,3]}?take=a&take=0'})
+        self.st.set_config(
+            {
+                "test_config": (
+                    'take://json://{"a":{"aa":0,"ab":1},"b":{"ba":2,"bb":3}}?take=b&take=ba'
+                )
+            }
+        )
         p = self.st.get_single_plugin(nt_test_run_id, "test_data")
-        self.assertEqual(p.test_config, 1)
+        self.assertEqual(p.test_config, 2)
 
     @unittest.skipIf(not straxen.utilix_is_configured(), "No db access, cannot test!")
     def test_bodedga_get(self):
@@ -434,9 +441,9 @@ def test_regex_url_warnings(self):
 
     def test_pad_array(self):
         """Test that pad_array works as expected."""
-
+        n = 3
         self.st.set_config(
-            {"test_config": "pad-array://json://[1,2,3]?pad_left=2&pad_right=3&pad_value=0"}
+            {"test_config": f"pad-array://range://{n}?pad_left=2&pad_right=3&pad_value=0"}
         )
         p = self.st.get_single_plugin(nt_test_run_id, "test_data")
         self.assertEqual(len(p.test_config), 8)