2279036581

Build Type

github

Committed by GitHub

Commit Message

<a href="https://github.com/XENONnT/straxen/commit/<a class=hub.com/XENONnT/straxen/commit/<a class="double-link" href="https://git"><a class=hub.com/XENONnT/straxen/commit/<a class="double-link" href="https://git"><a class=hub.com/XENONnT/straxen/commit/<a class="double-link" href="https://git"><a class=hub.com/XENONnT/straxen/commit/76f8f56ccb92bf8ccc4b5d4f2f432102a09d8e18">76f8f56cc">&lt;a href=&quot;https://github.com/XENONnT/straxen/commit/</a><a class="double-link" href="https://github.com/XENONnT/straxen/commit/&lt;a class=&quot;double-link&quot; href=&quot;https://git">&lt;a class=</a>hub.com/XENONnT/straxen/commit/&lt;a class=&quot;double-link&quot; href=&quot;https://git">&lt;a class=</a>hub.com/XENONnT/straxen/commit/&lt;a class=&quot;double-link&quot; href=&quot;https://git">&lt;a class=</a>hub.com/XENONnT/straxen/commit/&lt;a class=&quot;double-link&quot; href=&quot;https://git">&lt;a class=</a>hub.com/XENONnT/straxen/commit/76f8f56ccb92bf8ccc4b5d4f2f432102a09d8e18">76f8f56cc</a><a href="https://github.com/XENONnT/straxen/commit/76f8f56ccb92bf8ccc4b5d4f2f432102a09d8e18">&quot;&gt;&amp;lt;a href=&amp;quot;https://github.com/XENONnT/straxen/commit/&lt;/a&gt;&lt;a class=&quot;double-link&quot; href=&quot;https://github.com/XENONnT/straxen/commit/&amp;lt;a class=&amp;quot;double-link&amp;quot; href=&amp;quot;https://git&quot;&gt;&amp;lt;a class=&lt;/a&gt;hub.com/XENONnT/straxen/commit/&amp;lt;a class=&amp;quot;double-link&amp;quot; href=&amp;quot;https://git&quot;&gt;&amp;lt;a class=&lt;/a&gt;hub.com/XENONnT/straxen/commit/&amp;lt;a class=&amp;quot;double-link&amp;quot; href=&amp;quot;https://git&quot;&gt;&amp;lt;a class=&lt;/a&gt;hub.com/XENONnT/straxen/commit/&amp;lt;a class=&amp;quot;double-link&amp;quot; href=&amp;quot;https://git&quot;&gt;&amp;lt;a class=&lt;/a&gt;hub.com/XENONnT/straxen/commit/76f8f56ccb92bf8ccc4b5d4f2f432102a09d8e18&quot;&gt;76f8f56cc&lt;/a&gt;&lt;a href=&quot;https://github.com/XENONnT/straxen/commit/76f8f56ccb92bf8ccc4b5d4f2f432102a09d8e18&quot;&gt;&amp;quot;&amp;gt;&amp;amp;lt;a href=&amp;amp;quot;https://github.com/XENONnT/straxen/commit/&amp;lt;/a&amp;gt;&amp;lt;a class=&amp;quot;double-link&amp;quot; href=&amp;quot;https://github.com/XENONnT/straxen/commit/&amp;amp;lt;a class=&amp;amp;quot;double-link&am... (continued)

Pull Request Pull Request #1031: tweaks to #1027 - single more general protocol

Run Details

21 of 21 new or added lines in 2 files covered. (100.0%)

6826 of 7272 relevant lines covered (93.87%)

1.88 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

95.57

/straxen/url_config.py

import json
import typing
from typing import Container
import numpy as np
import strax
import fsspec
import pandas as pd
import straxen
import inspect
from urllib.parse import urlparse, parse_qs
from ast import literal_eval
from strax.config import OMITTED
import os
import tempfile
import tarfile
import pytz
from utilix import xent_collection
from scipy.interpolate import interp1d

export, __all__ = strax.exporter()

_CACHES = {}


@export
def clear_config_caches():
    for cache in _CACHES.values():
        cache.clear()


@export
def config_cache_size_mb():
    return straxen.total_size(_CACHES)//1e6


def parse_val(val: str):
    '''Attempt to parse a string value as
    a python literal, falls back to returning just
    the original string if cant be parsed.
    '''
    try:
        val = literal_eval(val)
    except ValueError:
        pass
    except SyntaxError:
        pass
    return val


@export
class URLConfig(strax.Config):
    """Dispatch on URL protocol.
    unrecognized protocol returns identity
    inspired by dasks Dispatch and fsspec fs protocols.
    """
    _LOOKUP = {}
    SCHEME_SEP = '://'
    QUERY_SEP = '?'
    PLUGIN_ATTR_PREFIX = 'plugin.'

    def __init__(self, cache=0, **kwargs):
        """
        :param cache: number of values to keep in cache, 
                      if set to True will cache all values
        :param **kwargs: additional keyword arguments accepted by strax.Option
        """
        self.final_type = OMITTED
        super().__init__(**kwargs)
        # Ensure backwards compatibility with Option validation
        # type of the config value can be different from the fetched value.
        if self.type is not OMITTED:
            self.final_type = self.type
            self.type = OMITTED  # do not enforce type on the URL
        if cache:
            cache_len = 100 if cache is True else int(cache) 
            cache = straxen.CacheDict(cache_len=cache_len)
            _CACHES[id(self)] = cache

    @property
    def cache(self):
        return _CACHES.get(id(self), {})

    @classmethod
    def register(cls, protocol, func=None):
        """Register dispatch of `func` on urls
        starting with protocol name `protocol` """

        def wrapper(func):
            if isinstance(protocol, tuple):
                for t in protocol:
                    cls.register(t, func)
                return func

            if not isinstance(protocol, str):
                raise ValueError('Protocol name must be a string.')

            if protocol in cls._LOOKUP:
                raise ValueError(f'Protocol with name {protocol} already registered.')
            cls._LOOKUP[protocol] = func
            return func
        return wrapper(func) if func is not None else wrapper

    def dispatch(self, url, *args, **kwargs):
        """
        Call the corresponding method based on protocol in url.
        chained protocols will be called with the result of the
        previous protocol as input
        overrides are passed to any protocol whos signature can accept them.
        """

        # separate the protocol name from the path
        protocol, _, path = url.partition(self.SCHEME_SEP)

        # find the corresponding protocol method
        meth = self._LOOKUP.get(protocol, None)
        if meth is None:
            # unrecognized protocol
            # evaluate as string-literal
            return url

        if self.SCHEME_SEP in path:
            # url contains a nested protocol
            # first call sub-protocol
            arg = self.dispatch(path, **kwargs)
        else:
            # we are at the end of the chain
            # method should be called with path as argument
            arg = path

        # filter kwargs to pass only the kwargs
        #  accepted by the method.
        kwargs = straxen.filter_kwargs(meth, kwargs)

        return meth(arg, *args, **kwargs)

    @classmethod
    def split_url_kwargs(cls, url):
        """split a url into path and kwargs
        """
        path, _, _ = url.partition(cls.QUERY_SEP)
        kwargs = {}
        for k, v in parse_qs(urlparse(url).query).items():
            # values of query arguments are evaluated as lists
            # split logic depending on length
            n = len(v)
            if not n:
                kwargs[k] = None
            elif n == 1:
                kwargs[k] = parse_val(v[0])
            else:
                kwargs[k] = list(map(parse_val, v))
        return path, kwargs

    @classmethod
    def kwarg_from_url(cls, url: str, key: str):
        path, kwargs = cls.split_url_kwargs(url)
        return kwargs.get(key, None)
        
    @classmethod
    def format_url_kwargs(cls, url, **kwargs):
        '''Add keyword arguments to a URL.
        Sorts all arguments by key for hash consistency
        '''
        url, extra_kwargs = cls.split_url_kwargs(url)
        kwargs = dict(extra_kwargs, **kwargs)
        arg_list = []
        for k, v in sorted(kwargs.items()):
            if isinstance(v, list):
                # lists are passed as multiple arguments with the same key
                arg_list.extend([f"{k}={vi}" for vi in v])
            else:
                arg_list.append(f"{k}={v}")
        arg_str = "&".join(arg_list)
        arg_str = cls.QUERY_SEP + arg_str if arg_str else ''
        return url + arg_str

    def fetch_attribute(self, plugin, value):
        if isinstance(value, str) and value.startswith(self.PLUGIN_ATTR_PREFIX):
            # kwarg is referring to a plugin attribute, lets fetch it
            return getattr(plugin, value[len(self.PLUGIN_ATTR_PREFIX):], value)

        if isinstance(value, list):
            return [self.fetch_attribute(plugin, v) for v in value]

        # kwarg is a literal, add its value to the kwargs dict
        return value

    def fetch(self, plugin):
        """override the Config.fetch method
        this is called when the attribute is accessed
        """
        # first fetch the user-set value 
        # from the config dictionary
        url = super().fetch(plugin)

        if not isinstance(url, str):
            # if the value is not a string it is evaluated
            # as a literal config and returned as is.
            return url

        if self.SCHEME_SEP not in url:
            # no protocol in the url so its evaluated 
            # as string-literal config and returned as is
            return url

        # separate out the query part of the URL which
        # will become the method kwargs
        url, url_kwargs = self.split_url_kwargs(url)

        kwargs = {k: self.fetch_attribute(plugin, v)
                  for k, v in url_kwargs.items()}

        # construct a deterministic hash key
        key = strax.deterministic_hash((url, kwargs))

        # fetch from cache if exists
        value = self.cache.get(key, None)

        # not in cache, lets fetch it
        if value is None:
            value = self.dispatch(url, **kwargs)
            self.cache[key] = value

        return value

    @classmethod
    def protocol_descr(cls):
        rows = []
        for k, v in cls._LOOKUP.items():
            row = {
                'name': f"{k}://",
                'description': v.__doc__,
                'signature': str(inspect.signature(v)),
            }
            rows.append(row)
        return pd.DataFrame(rows)

    @classmethod
    def print_protocols(cls):
        df = cls.protocol_descr()
        print(df)


@URLConfig.register('cmt')
def get_correction(name: str,
                   run_id: str = None,
                   version: str = 'ONLINE',
                   detector: str = 'nt',
                   **kwargs):
    """Get value for name from CMT"""
    if run_id is None:
        raise ValueError('Attempting to fetch a correction without a run id.')
    return straxen.get_correction_from_cmt(run_id, (name, version, detector == 'nt'))


@URLConfig.register('resource')
def get_resource(name: str,
                 fmt: str = 'text',
                 **kwargs):
    """
    Fetch a straxen resource

    Allow a direct download using <fmt='abs_path'> otherwise kwargs are
    passed directly to straxen.get_resource.
    """
    if fmt == 'abs_path':
        downloader = straxen.MongoDownloader()
        return downloader.download_single(name)
    return straxen.get_resource(name, fmt=fmt)


@URLConfig.register('fsspec')
def read_file(path: str, **kwargs):
    """Support fetching files from arbitrary filesystems
    """
    with fsspec.open(path, **kwargs) as f:
        content = f.read()
    return content


@URLConfig.register('json')
def read_json(content: str, **kwargs):
    """Load json string as a python object
    """
    return json.loads(content)


@URLConfig.register('take')
def get_key(container: Container, take=None, **kwargs):
    """ return a single element of a container
    """
    if take is None:
        return container
    if not isinstance(take, list):
        take = [take]

    # support for multiple keys for
    # nested objects
    for t in take:
        container = container[t]

    return container


@URLConfig.register('format')
def format_arg(arg: str, **kwargs):
    """apply pythons builtin format function to a string"""
    return arg.format(**kwargs)


@URLConfig.register('itp_map')
def load_map(some_map, method='WeightedNearestNeighbors', **kwargs):
    """Make an InterpolatingMap"""
    return straxen.InterpolatingMap(some_map, method=method, **kwargs)


@URLConfig.register('bodega')
def load_value(name: str, bodega_version=None):
    """Load a number from BODEGA file"""
    if bodega_version is None:
        raise ValueError('Provide version see e.g. tests/test_url_config.py')
    nt_numbers = straxen.get_resource("XENONnT_numbers.json", fmt="json")
    return nt_numbers[name][bodega_version]["value"]


@URLConfig.register('tf')
def open_neural_net(model_path: str, **kwargs):
    # Nested import to reduce loading time of import straxen and it not
    # base requirement
    import tensorflow as tf
    if not os.path.exists(model_path):
        raise FileNotFoundError(f'No file at {model_path}')
    with tempfile.TemporaryDirectory() as tmpdirname:
        tar = tarfile.open(model_path, mode="r:gz")
        tar.extractall(path=tmpdirname)
        return tf.keras.models.load_model(tmpdirname)


@URLConfig.register('itp_dict')
def get_itp_dict(loaded_json,
                 run_id=None,
                 time_key='time',
                 itp_keys='correction',
                 **kwargs) -> typing.Union[np.ndarray, typing.Dict[str, np.ndarray]]:
    """
    Interpolate the dictionary at the start time that is queried from
    a run-id.

    The loaded JSON (dictionary), should have at least one collum `time`
    and has other columns given by `itp_dict_keys`.

    :param loaded_json: a dictionary with a time-series
    :param run_id: run_id
    :param time_key: key that gives the timestamps
    :param itp_keys: which keys from the dict to read. Should be
        comma (',') separated!

    :return: Interpolated values of dict at the start time, either
        returned as an np.ndarray (single value) or as a dict
        (multiple itp_dict_keys)
    """
    keys = strax.to_str_tuple(itp_keys.split(','))
    times = loaded_json[time_key]
    if not all(k in loaded_json for k in keys):
        raise ValueError(f'One or more of {itp_keys} are not in {loaded_json.keys()}')

    # get start time of this run. Need to make tz-aware
    start = xent_collection().find_one({'number': int(run_id)}, {'start': 1})['start']
    start = pytz.utc.localize(start).timestamp() * 1e9

    try:
        if len(strax.to_str_tuple(keys))>1:
            return {key:
                    interp1d(times, loaded_json[key], bounds_error=True)(start)
                    for key in keys}

        else:
            interp = interp1d(times, loaded_json[keys[0]], bounds_error=True)
            return interp(start)
    except ValueError as e:
        raise ValueError(f"The correction is not defined for run {run_id}") from e


@URLConfig.register('rekey_dict')
def rekey_dict(d, replace_keys='', with_keys=''):
    '''
    :param d: dictionary that will have its keys renamed
    :param replace_keys: comma-separated list of keys that will be replaced
    :param with_keys:  comma-separated list of keys that will replace the replace_keys
    :return: dictionary with renamed keys
    '''
    new_dict = d.copy()
    replace_keys = strax.to_str_tuple(replace_keys.split(','))
    with_keys = strax.to_str_tuple(with_keys.split(','))
    if len(replace_keys) != len(with_keys):
        raise RuntimeError("replace_keys and with_keys must have the same length")
    for old_key, new_key in zip(replace_keys, with_keys):
        new_dict[new_key] = new_dict.pop(old_key)
    return new_dict

1	import json	2✔
2	import typing	2✔
3	from typing import Container	2✔
4	import numpy as np	2✔
5	import strax	2✔
6	import fsspec	2✔
7	import pandas as pd	2✔
8	import straxen	2✔
9	import inspect	2✔
10	from urllib.parse import urlparse, parse_qs	2✔
11	from ast import literal_eval	2✔
12	from strax.config import OMITTED	2✔
13	import os	2✔
14	import tempfile	2✔
15	import tarfile	2✔
16	import pytz	2✔
17	from utilix import xent_collection	2✔
18	from scipy.interpolate import interp1d	2✔
19
20	export, __all__ = strax.exporter()	2✔
21
22	_CACHES = {}	2✔
23
24
25	@export	2✔
26	def clear_config_caches():	2✔
27	for cache in _CACHES.values():	2✔
28	cache.clear()	2✔
29
30
31	@export	2✔
32	def config_cache_size_mb():	2✔
33	return straxen.total_size(_CACHES)//1e6	2✔
34
35
36	def parse_val(val: str):	2✔
37	'''Attempt to parse a string value as
38	a python literal, falls back to returning just
39	the original string if cant be parsed.
40	'''
41	try:	2✔
42	val = literal_eval(val)	2✔
43	except ValueError:	2✔
44	pass	2✔
45	except SyntaxError:	2✔
46	pass	2✔
47	return val	2✔
48
49
50	@export	2✔
51	class URLConfig(strax.Config):	2✔
52	"""Dispatch on URL protocol.
53	unrecognized protocol returns identity
54	inspired by dasks Dispatch and fsspec fs protocols.
55	"""
56	_LOOKUP = {}	2✔
57	SCHEME_SEP = '://'	2✔
58	QUERY_SEP = '?'	2✔
59	PLUGIN_ATTR_PREFIX = 'plugin.'	2✔
60
61	def __init__(self, cache=0, **kwargs):	2✔
62	"""
63	:param cache: number of values to keep in cache,
64	if set to True will cache all values
65	:param **kwargs: additional keyword arguments accepted by strax.Option
66	"""
67	self.final_type = OMITTED	2✔
68	super().__init__(**kwargs)	2✔
69	# Ensure backwards compatibility with Option validation
70	# type of the config value can be different from the fetched value.
71	if self.type is not OMITTED:	2✔
72	self.final_type = self.type	2✔
73	self.type = OMITTED # do not enforce type on the URL	2✔
74	if cache:	2✔
75	cache_len = 100 if cache is True else int(cache)	2✔
76	cache = straxen.CacheDict(cache_len=cache_len)	2✔
77	_CACHES[id(self)] = cache	2✔
78
79	@property	2✔
80	def cache(self):	2✔
81	return _CACHES.get(id(self), {})	2✔
82
83	@classmethod	2✔
84	def register(cls, protocol, func=None):	2✔
85	"""Register dispatch of `func` on urls
86	starting with protocol name `protocol` """
87
88	def wrapper(func):	2✔
89	if isinstance(protocol, tuple):	2✔
90	for t in protocol:	×
91	cls.register(t, func)	×
92	return func	×
93
94	if not isinstance(protocol, str):	2✔
95	raise ValueError('Protocol name must be a string.')
96
97	if protocol in cls._LOOKUP:	2✔
98	raise ValueError(f'Protocol with name {protocol} already registered.')
99	cls._LOOKUP[protocol] = func	2✔
100	return func	2✔
101	return wrapper(func) if func is not None else wrapper	2✔
102
103	def dispatch(self, url, args, *kwargs):	2✔
104	"""
105	Call the corresponding method based on protocol in url.
106	chained protocols will be called with the result of the
107	previous protocol as input
108	overrides are passed to any protocol whos signature can accept them.
109	"""
110
111	# separate the protocol name from the path
112	protocol, _, path = url.partition(self.SCHEME_SEP)	2✔
113
114	# find the corresponding protocol method
115	meth = self._LOOKUP.get(protocol, None)	2✔
116	if meth is None:	2✔
117	# unrecognized protocol
118	# evaluate as string-literal
119	return url	2✔
120
121	if self.SCHEME_SEP in path:	2✔
122	# url contains a nested protocol
123	# first call sub-protocol
124	arg = self.dispatch(path, **kwargs)	2✔
125	else:
126	# we are at the end of the chain
127	# method should be called with path as argument
128	arg = path	2✔
129
130	# filter kwargs to pass only the kwargs
131	# accepted by the method.
132	kwargs = straxen.filter_kwargs(meth, kwargs)	2✔
133
134	return meth(arg, args, *kwargs)	2✔
135
136	@classmethod	2✔
137	def split_url_kwargs(cls, url):	2✔
138	"""split a url into path and kwargs
139	"""
140	path, _, _ = url.partition(cls.QUERY_SEP)	2✔
141	kwargs = {}	2✔
142	for k, v in parse_qs(urlparse(url).query).items():	2✔
143	# values of query arguments are evaluated as lists
144	# split logic depending on length
145	n = len(v)	2✔
146	if not n:	2✔
147	kwargs[k] = None	×
148	elif n == 1:	2✔
149	kwargs[k] = parse_val(v[0])	2✔
150	else:
151	kwargs[k] = list(map(parse_val, v))	2✔
152	return path, kwargs	2✔
153
154	@classmethod	2✔
155	def kwarg_from_url(cls, url: str, key: str):	2✔
156	path, kwargs = cls.split_url_kwargs(url)	2✔
157	return kwargs.get(key, None)	2✔
158
159	@classmethod	2✔
160	def format_url_kwargs(cls, url, **kwargs):	2✔
161	'''Add keyword arguments to a URL.
162	Sorts all arguments by key for hash consistency
163	'''
164	url, extra_kwargs = cls.split_url_kwargs(url)	2✔
165	kwargs = dict(extra_kwargs, **kwargs)	2✔
166	arg_list = []	2✔
167	for k, v in sorted(kwargs.items()):	2✔
168	if isinstance(v, list):	2✔
169	# lists are passed as multiple arguments with the same key
170	arg_list.extend([f"{k}={vi}" for vi in v])	×
171	else:
172	arg_list.append(f"{k}={v}")	2✔
173	arg_str = "&".join(arg_list)	2✔
174	arg_str = cls.QUERY_SEP + arg_str if arg_str else ''	2✔
175	return url + arg_str	2✔
176
177	def fetch_attribute(self, plugin, value):	2✔
178	if isinstance(value, str) and value.startswith(self.PLUGIN_ATTR_PREFIX):	2✔
179	# kwarg is referring to a plugin attribute, lets fetch it
180	return getattr(plugin, value[len(self.PLUGIN_ATTR_PREFIX):], value)	2✔
181
182	if isinstance(value, list):	2✔
183	return [self.fetch_attribute(plugin, v) for v in value]	2✔
184
185	# kwarg is a literal, add its value to the kwargs dict
186	return value	2✔
187
188	def fetch(self, plugin):	2✔
189	"""override the Config.fetch method
190	this is called when the attribute is accessed
191	"""
192	# first fetch the user-set value
193	# from the config dictionary
194	url = super().fetch(plugin)	2✔
195
196	if not isinstance(url, str):	2✔
197	# if the value is not a string it is evaluated
198	# as a literal config and returned as is.
199	return url	2✔
200
201	if self.SCHEME_SEP not in url:	2✔
202	# no protocol in the url so its evaluated
203	# as string-literal config and returned as is
204	return url	2✔
205
206	# separate out the query part of the URL which
207	# will become the method kwargs
208	url, url_kwargs = self.split_url_kwargs(url)	2✔
209
210	kwargs = {k: self.fetch_attribute(plugin, v)	2✔
211	for k, v in url_kwargs.items()}
212
213	# construct a deterministic hash key
214	key = strax.deterministic_hash((url, kwargs))	2✔
215
216	# fetch from cache if exists
217	value = self.cache.get(key, None)	2✔
218
219	# not in cache, lets fetch it
220	if value is None:	2✔
221	value = self.dispatch(url, **kwargs)	2✔
222	self.cache[key] = value	2✔
223
224	return value	2✔
225
226	@classmethod	2✔
227	def protocol_descr(cls):	2✔
228	rows = []	2✔
229	for k, v in cls._LOOKUP.items():	2✔
230	row = {	2✔
231	'name': f"{k}://",
232	'description': v.__doc__,
233	'signature': str(inspect.signature(v)),
234	}
235	rows.append(row)	2✔
236	return pd.DataFrame(rows)	2✔
237
238	@classmethod	2✔
239	def print_protocols(cls):	2✔
240	df = cls.protocol_descr()	2✔
241	print(df)	2✔
242
243
244	@URLConfig.register('cmt')	2✔
245	def get_correction(name: str,	2✔
246	run_id: str = None,
247	version: str = 'ONLINE',
248	detector: str = 'nt',
249	**kwargs):
250	"""Get value for name from CMT"""
251	if run_id is None:	2✔
252	raise ValueError('Attempting to fetch a correction without a run id.')
253	return straxen.get_correction_from_cmt(run_id, (name, version, detector == 'nt'))	2✔
254
255
256	@URLConfig.register('resource')	2✔
257	def get_resource(name: str,	2✔
258	fmt: str = 'text',
259	**kwargs):
260	"""
261	Fetch a straxen resource
262
263	Allow a direct download using <fmt='abs_path'> otherwise kwargs are
264	passed directly to straxen.get_resource.
265	"""
266	if fmt == 'abs_path':	2✔
267	downloader = straxen.MongoDownloader()	2✔
268	return downloader.download_single(name)	2✔
269	return straxen.get_resource(name, fmt=fmt)	2✔
270
271
272	@URLConfig.register('fsspec')	2✔
273	def read_file(path: str, **kwargs):	2✔
274	"""Support fetching files from arbitrary filesystems
275	"""
276	with fsspec.open(path, **kwargs) as f:	2✔
277	content = f.read()	2✔
278	return content	2✔
279
280
281	@URLConfig.register('json')	2✔
282	def read_json(content: str, **kwargs):	2✔
283	"""Load json string as a python object
284	"""
285	return json.loads(content)	2✔
286
287
288	@URLConfig.register('take')	2✔
289	def get_key(container: Container, take=None, **kwargs):	2✔
290	""" return a single element of a container
291	"""
292	if take is None:	2✔
293	return container	×
294	if not isinstance(take, list):	2✔
295	take = [take]	2✔
296
297	# support for multiple keys for
298	# nested objects
299	for t in take:	2✔
300	container = container[t]	2✔
301
302	return container	2✔
303
304
305	@URLConfig.register('format')	2✔
306	def format_arg(arg: str, **kwargs):	2✔
307	"""apply pythons builtin format function to a string"""
308	return arg.format(**kwargs)	2✔
309
310
311	@URLConfig.register('itp_map')	2✔
312	def load_map(some_map, method='WeightedNearestNeighbors', **kwargs):	2✔
313	"""Make an InterpolatingMap"""
314	return straxen.InterpolatingMap(some_map, method=method, **kwargs)	2✔
315
316
317	@URLConfig.register('bodega')	2✔
318	def load_value(name: str, bodega_version=None):	2✔
319	"""Load a number from BODEGA file"""
320	if bodega_version is None:	2✔
321	raise ValueError('Provide version see e.g. tests/test_url_config.py')
322	nt_numbers = straxen.get_resource("XENONnT_numbers.json", fmt="json")	2✔
323	return nt_numbers[name][bodega_version]["value"]	2✔
324
325
326	@URLConfig.register('tf')	2✔
327	def open_neural_net(model_path: str, **kwargs):	2✔
328	# Nested import to reduce loading time of import straxen and it not
329	# base requirement
330	import tensorflow as tf	2✔
331	if not os.path.exists(model_path):	2✔
332	raise FileNotFoundError(f'No file at {model_path}')
333	with tempfile.TemporaryDirectory() as tmpdirname:	2✔
334	tar = tarfile.open(model_path, mode="r:gz")	2✔
335	tar.extractall(path=tmpdirname)	2✔
336	return tf.keras.models.load_model(tmpdirname)	2✔
337
338
339	@URLConfig.register('itp_dict')	2✔
340	def get_itp_dict(loaded_json,	2✔
341	run_id=None,
342	time_key='time',
343	itp_keys='correction',
344	**kwargs) -> typing.Union[np.ndarray, typing.Dict[str, np.ndarray]]:
345	"""
346	Interpolate the dictionary at the start time that is queried from
347	a run-id.
348
349	The loaded JSON (dictionary), should have at least one collum `time`
350	and has other columns given by `itp_dict_keys`.
351
352	:param loaded_json: a dictionary with a time-series
353	:param run_id: run_id
354	:param time_key: key that gives the timestamps
355	:param itp_keys: which keys from the dict to read. Should be
356	comma (',') separated!
357
358	:return: Interpolated values of dict at the start time, either
359	returned as an np.ndarray (single value) or as a dict
360	(multiple itp_dict_keys)
361	"""
362	keys = strax.to_str_tuple(itp_keys.split(','))	2✔
363	times = loaded_json[time_key]	2✔
364	if not all(k in loaded_json for k in keys):	2✔
365	raise ValueError(f'One or more of {itp_keys} are not in {loaded_json.keys()}')
366
367	# get start time of this run. Need to make tz-aware
368	start = xent_collection().find_one({'number': int(run_id)}, {'start': 1})['start']	2✔
369	start = pytz.utc.localize(start).timestamp() * 1e9	2✔
370
371	try:	2✔
372	if len(strax.to_str_tuple(keys))>1:	2✔
373	return {key:	2✔
374	interp1d(times, loaded_json[key], bounds_error=True)(start)
375	for key in keys}
376
377	else:
378	interp = interp1d(times, loaded_json[keys[0]], bounds_error=True)	×
379	return interp(start)	×
380	except ValueError as e:	×
381	raise ValueError(f"The correction is not defined for run {run_id}") from e
382
383
384	@URLConfig.register('rekey_dict')	2✔
385	def rekey_dict(d, replace_keys='', with_keys=''):	2✔
386	'''
387	:param d: dictionary that will have its keys renamed
388	:param replace_keys: comma-separated list of keys that will be replaced
389	:param with_keys: comma-separated list of keys that will replace the replace_keys
390	:return: dictionary with renamed keys
391	'''
392	new_dict = d.copy()	2✔
393	replace_keys = strax.to_str_tuple(replace_keys.split(','))	2✔
394	with_keys = strax.to_str_tuple(with_keys.split(','))	2✔
395	if len(replace_keys) != len(with_keys):	2✔
396	raise RuntimeError("replace_keys and with_keys must have the same length")
397	for old_key, new_key in zip(replace_keys, with_keys):	2✔
398	new_dict[new_key] = new_dict.pop(old_key)	2✔
399	return new_dict	2✔

XENONnT / straxen / 2279036581

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous