New URLConfig protocols - list-to-array and list-to-dict (#1104)

* Adding new protocols to change the format of the values from URLConfigurations

* Second comit, fixed indentation

* fix up a bit

* add tests

* trailing newline

* Update tests/test_url_config.py

Co-authored-by: Joran R. Angevaare <joranangevaare@gmail.com>

* jorans comments

Co-authored-by: Yossi Mosbacher <joe.mosbacher>
Co-authored-by: Joran R. Angevaare <jorana@nikhef.nl>
Co-authored-by: Yossi Mosbacher <joe.mosbacher@gmail.com>
Co-authored-by: Joran R. Angevaare <joranangevaare@gmail.com>

straxen/url_config.py

@@ -1,6 +1,6 @@
 import json
 import typing
-from typing import Container
+from typing import Container, Iterable
 from uuid import uuid4
 import numpy as np
 import strax
@@ -9,6 +9,8 @@
 import straxen
 import inspect
 from urllib.parse import urlparse, parse_qs
+from immutabledict import immutabledict
+
 from ast import literal_eval
 from strax.config import OMITTED
 import os
@@ -48,6 +50,12 @@ def parse_val(val: str):
     return val
 
 
+def get_item_or_attr(obj, key, default=None):
+    if isinstance(obj, dict):
+        return obj.get(key, default)
+    return getattr(obj, key, default)
+
+
 @export
 class URLConfig(strax.Config):
     """Dispatch on URL protocol.
@@ -448,3 +456,40 @@ def rekey_dict(d, replace_keys='', with_keys=''):
     for old_key, new_key in zip(replace_keys, with_keys):
         new_dict[new_key] = new_dict.pop(old_key)
     return new_dict
+
+
+@URLConfig.register('objects-to-dict')
+def objects_to_dict(objects: list, key_attr=None, value_attr='value', immutable=False):
+    '''Converts a list of objects/dicts to a single dictionary by taking the 
+    key and value from each of the objects/dicts. If key_attr is not provided,
+    the list index is used as the key.
+
+    :param objects: list of objects/dicts that will be converted to a dictionary
+    :param key_attr: key/attribute of the objects that will be used as key in the dictionary
+    :param value_attr: key/attribute of the objects that will be used as value in the dictionary
+    '''
+    if not isinstance(objects, Iterable):
+        raise TypeError(f'The objects-to-dict protocol expects an iterable '
+                        f'of objects but received {type(objects)} instead.')
+    result = {}
+    for i, obj in enumerate(objects):
+        key = i if key_attr is None else get_item_or_attr(obj, key_attr)
+        result[key] = get_item_or_attr(obj, value_attr)
+
+    if immutable:
+        result = immutabledict(result)
+
+    return result
+
+
+@URLConfig.register('list-to-array')
+def objects_to_array(objects: list):
+    '''
+    Converts a list of objects/dicts to a numpy array.
+    :param objects: Any list of objects'''
+
+    if not isinstance(objects, Iterable):
+        raise TypeError(f'The list-to-array protocol expects an '
+                        f'iterable but recieved a {type(objects)} instead')
+
+    return np.array(objects)

tests/test_url_config.py

@@ -14,6 +14,11 @@
 from datetime import datetime
 
 
+class DummyObject:
+    def __init__(self, **kwargs):
+        self.__dict__.update(kwargs)
+
+
 @straxen.URLConfig.register('random')
 def generate_random(_):
     return random.random()
@@ -29,6 +34,12 @@ def large_array(_):
     return np.ones(1_000_000).tolist()
 
 
+@straxen.URLConfig.register('object-list')
+def object_list(length):
+    length = int(length)
+    return [DummyObject(a=i, b=i+1) for i in range(length)]
+
+
 class ExamplePlugin(strax.Plugin):
     depends_on = ()
     dtype = strax.time_fields
@@ -270,3 +281,15 @@ def test_dry_evaluation(self):
         # as in the plugin_url and should complain about that
         with self.assertRaises(ValueError):
             straxen.URLConfig.evaluate_dry(plugin_url)
+
+    def test_objects_to_dict(self):
+        n = 3
+        self.st.set_config({'test_config': f'objects-to-dict://object-list://{n}?key_attr=a&value_attr=b'})
+        p = self.st.get_single_plugin(nt_test_run_id, 'test_data')
+        self.assertEqual(p.test_config, {i:i+1 for i in range(n)})
+
+    def test_list_to_array(self):
+        n = 3
+        self.st.set_config({'test_config': f'list-to-array://object-list://{n}'})
+        p = self.st.get_single_plugin(nt_test_run_id, 'test_data')
+        self.assertIsInstance(p.test_config, np.ndarray)