diff --git a/docs/source/reference/strax.rst b/docs/source/reference/strax.rst index 55d42899..bf62a60b 100644 --- a/docs/source/reference/strax.rst +++ b/docs/source/reference/strax.rst @@ -93,6 +93,14 @@ strax.testutils module :undoc-members: :show-inheritance: +strax.url\_config module +------------------------ + +.. automodule:: strax.url_config + :members: + :undoc-members: + :show-inheritance: + strax.utils module ------------------ diff --git a/pyproject.toml b/pyproject.toml index 011c6c9b..63db7ab3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,6 +31,7 @@ dependencies = [ "numcodecs<0.16.0", "packaging", "pandas", + "pydantic>=1.10", "pymongo", "pytz", "scipy", diff --git a/strax/__init__.py b/strax/__init__.py index 2de548da..129961a9 100644 --- a/strax/__init__.py +++ b/strax/__init__.py @@ -17,6 +17,7 @@ from .storage.zipfiles import * from .config import * +from .url_config import * from .plugins import * from .mailbox import * diff --git a/strax/url_config.py b/strax/url_config.py new file mode 100644 index 00000000..84293312 --- /dev/null +++ b/strax/url_config.py @@ -0,0 +1,519 @@ +import json +import strax + +import numbers +import inspect +from ast import literal_eval +from typing import Any, Dict, Optional, Mapping, Union + +import pandas as pd +from urllib.parse import urlparse, parse_qs + +try: + # pydantic >= 2 ships the v1 API under pydantic.v1 + from pydantic.v1.validators import find_validators + from pydantic.v1.config import get_config +except ImportError: + from pydantic.validators import find_validators + from pydantic.config import get_config + +from strax.config import OMITTED +from strax.utils import filter_kwargs, CacheDict, total_size + +export, __all__ = strax.exporter() + +_CACHES: Dict[int, Any] = {} + + +@export +def clear_config_caches(): + for cache in _CACHES.values(): + cache.clear() + + +@export +def config_cache_size_mb(): + return 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: Dict[str, Any] = {} + _PREPROCESSORS = () + + SCHEME_SEP = "://" + QUERY_SEP = "?" + NAMESPACE_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 = 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 + + @classmethod + def preprocessor(cls, func=None, precedence=0): + """Register a new processor to modify the config values before they are used.""" + + def wrapper(func): + entry = (precedence, func) + if entry in cls._PREPROCESSORS: + raise ValueError(f"This processor is already registered.") + cls._PREPROCESSORS += (entry,) + return func + + return wrapper(func) if func is not None else wrapper + + @classmethod + def eval( + cls, protocol: str, arg: Optional[Union[str, tuple]] = None, kwargs: Optional[dict] = None + ): + """Evaluate a URL/AST by recusively dispatching protocols by name with argument arg and + keyword arguments kwargs and return the value. + + If protocol does not exist, returnes arg + :param protocol: name of the protocol or a URL + :param arg: argument to pass to protocol, can be another (sub-protocol, arg, kwargs) tuple, + in which case sub-protocol will be evaluated and passed to protocol + :param kwargs: keyword arguments to be passed to the protocol + :return: (Any) The return value of the protocol on these arguments + + """ + + if protocol is not None and arg is None: + protocol, arg, kwargs = cls.url_to_ast(protocol) + + if protocol is None: + return arg + + if kwargs is None: + kwargs = {} + + meth = cls._LOOKUP[protocol] + + if isinstance(arg, tuple): + arg = cls.eval(*arg) + + # Just to be on the safe side + kwargs = filter_kwargs(meth, kwargs) + + return meth(arg, **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 + + @classmethod + def lookup_value(cls, value, **namespace): + """Optionally fetch an attribute from namespace if value is a string with cls.NAMESPACE_SEP + in it, the string is split and the first part is used to lookup an object in namespace and + the second part is used to lookup the value in the object. + + If the value is not a string or the target object is not in the namesapce, the value is + returned as is. + + """ + + if isinstance(value, list): + return [cls.lookup_value(v, **namespace) for v in value] + + if isinstance(value, str) and cls.NAMESPACE_SEP in value: + name, _, key = value.partition(cls.NAMESPACE_SEP) + if name in namespace: + obj = namespace[name] + if isinstance(obj, Mapping): + value = obj.get(key, value) + else: + value = getattr(obj, key, value) + + return value + + @classmethod + def deref_ast(cls, protocol, arg, kwargs, **namespace): + """Dereference an AST by looking up values in namespace.""" + if isinstance(arg, tuple): + arg = cls.deref_ast(*arg, **namespace) + else: + arg = cls.lookup_value(arg, **namespace) + kwargs = {k: cls.lookup_value(v, **namespace) for k, v in kwargs.items()} + return protocol, arg, kwargs + + def validate( + self, + config, + run_id=None, # TODO: will soon be removed + run_defaults=None, + set_defaults=True, + ): + """This method is called by the context on plugin initialization at this stage, the run_id + and context config are already known but the config values are not yet set on the plugin. + + Therefore its the perfect place to run any preprocessors on the config values to make any + needed changes before the configs are hashed. + + """ + super().validate(config, run_id, run_defaults, set_defaults) + + cfg = config[self.name] + + sorted_preprocessors = reversed(sorted(self._PREPROCESSORS, key=lambda x: x[0])) + + full_kwargs = dict( + name=self.name, + run_id=run_id, + run_defaults=run_defaults, + set_defaults=set_defaults, + ) + + for _, preprocessor in sorted_preprocessors: + kwargs = filter_kwargs(preprocessor, full_kwargs) + new_cfg = preprocessor(cfg, **kwargs) + cfg = new_cfg if new_cfg is not None else cfg + + config[self.name] = cfg + + if not isinstance(cfg, str) or self.SCHEME_SEP not in cfg: + # if the value is not a url config it is validated against + # its intended type (final_type) + self.validate_type(cfg) + + def validate_type(self, value): + """Validate the type of a value against its intended type.""" + msg = ( + f"Invalid type for option {self.name}. " + f"Expected a {self.final_type} instance, got {type(value)}" + ) + + if self.final_type is not OMITTED: + # Use pydantic to validate the type + # its validation is more flexible than isinstance + # it will coerce standard equivalent types + cfg = get_config(dict(arbitrary_types_allowed=True)) + if isinstance(self.final_type, tuple): + validators = [v for t in self.final_type for v in find_validators(t, config=cfg)] + else: + validators = find_validators(self.final_type, config=cfg) + for validator in validators: + try: + validator(value) + break + except Exception: + pass + else: + raise TypeError(msg) + + return value + + def fetch(self, plugin): + """Override the Config.fetch method this is called when the attribute is accessed from + withing the Plugin instance.""" + # 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 self.validate_type(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 self.validate_type(url) + + # evaluate the url as AST + protocol, arg, kwargs = self.url_to_ast(url) + + # allow run_id to be missing + run_id = getattr(plugin, "run_id", "000000") + + # construct a deterministic hash key from AST + key = strax.deterministic_hash((plugin.config, run_id, protocol, arg, kwargs)) + + # fetch from cache if exists + value = self.cache.get(key, None) + + # not in cache, lets fetch it + if value is None: + # resolve any referenced to plugin or config attributes + protocol, arg, kwargs = self.deref_ast( + protocol, arg, kwargs, config=plugin.config, plugin=plugin + ) + + value = self.eval(protocol, arg, kwargs) + value = self.validate_type(value) + self.cache[key] = value + + return value + + @classmethod + def ast_to_url( + cls, + protocol: Union[str, tuple], + arg: Optional[Union[str, tuple]] = None, + kwargs: Optional[dict] = None, + ): + """Convert a protocol abstract syntax tree to a valid URL.""" + + if isinstance(protocol, tuple): + protocol, arg, kwargs = protocol + + if kwargs is None: + kwargs = {} + + if protocol is None: + return arg + + if isinstance(arg, (list, dict, numbers.Number)) and protocol != "json": + arg = ( + "json", + json.dumps(arg), + ) + + if isinstance(arg, tuple): + arg = cls.ast_to_url(*arg) + + if not isinstance(arg, str): + raise TypeError(f"Type {type(arg)} is not supported as an argument.") + + arg, extra_kwargs = cls.split_url_kwargs(arg) + + kwargs.update(extra_kwargs) + + url = f"{protocol}{cls.SCHEME_SEP}{arg}" + + url = cls.format_url_kwargs(url, **kwargs) + + return url + + @classmethod + def url_to_ast(cls, url, **kwargs): + """Convert a URL to a protocol abstract syntax tree.""" + if not isinstance(url, str): + raise TypeError(f"URL must be a string, got {type(url)}") + + if cls.SCHEME_SEP not in url: + # no protocol in the url so its evaluated + # as string-literal config and returned as is + return None, url, {} + + # separate the protocol name from the path + protocol, _, path = url.partition(cls.SCHEME_SEP) + + # find the corresponding protocol method + meth = cls._LOOKUP.get(protocol, None) + if meth is None: + # unrecognized protocol + # evaluate as string-literal + return None, url, {} + + arg, url_kwargs = cls.split_url_kwargs(path) + kwargs.update(url_kwargs) + + if cls.SCHEME_SEP in arg: + # url contains a nested protocol + # first parsce sub-protocol + arg = cls.url_to_ast(arg, **kwargs) + + # Filter unused kwargs for this method. + # This is done also at the eval level but + # probably better to be safe. + kwargs = filter_kwargs(meth, kwargs) + + # Always sort kwargs for consistent ASTs + kwargs = dict(sorted(kwargs.items())) + + return protocol, arg, kwargs + + @classmethod + def are_equal(cls, first, second): + """Return whether two URLs are equivalent (have equal ASTs)""" + return cls.url_to_ast(first) == cls.url_to_ast(second) + + @classmethod + def protocol_descr(cls): + rows = [] + for k, v in cls._LOOKUP.items(): + descr = v.__doc__ + if descr is not None: + descr = descr.split("\n")[0] + + row = { + "name": f"{k}://", + "description": descr, + "signature": str(inspect.signature(v)), + "location": v.__module__, + } + rows.append(row) + return pd.DataFrame(rows) + + @classmethod + def print_protocols(cls): + df = cls.protocol_descr() + if len(df): + print(df) + else: + print("No protocols registered.") + + @classmethod + def preprocessor_descr(cls): + rows = [] + for k, v in cls._PREPROCESSORS: + descr = v.__doc__ + if descr is not None: + descr = descr.split("\n")[0] + row = { + "precedence": k, + "description": descr, + "signature": str(inspect.signature(v)), + "location": v.__module__, + } + rows.append(row) + return pd.DataFrame(rows) + + @classmethod + def print_preprocessors(cls): + df = cls.preprocessor_descr() + if len(df): + print(df) + else: + print("No Preprocessors registered.") + + @classmethod + def print_summary(cls): + print("=" * 30 + " Protocols " + "=" * 30) + cls.print_protocols() + print("=" * 30 + " Preprocessors " + "=" * 30) + cls.print_preprocessors() + + @classmethod + def evaluate_dry(cls, url: str, **kwargs): + """Utility function to quickly test and evaluate URL configs, without the initialization of + plugins (so no plugin attributes). plugin attributes can be passed as keyword arguments. + + example:: + + from strax import URLConfig + url_string='xedocs://electron_drift_velocity?run_id=027000&version=ONLINE' + URLConfig.evaluate_dry(url_string) + + # or similarly + url_string='xedocs://electron_drift_velocity?run_id=plugin.run_id&version=ONLINE' + URLConfig.evaluate_dry(url_string, run_id='027000') + + Please note that this has to be done outside of the plugin, so any + attributes of the plugin are not yet note to this dry evaluation + of the url-string. + + :param url: URL to evaluate, see above for example. + :keyword: any additional kwargs are passed to self.dispatch (see example) + :return: evaluated value of the URL. + + """ + url = cls.format_url_kwargs(url, **kwargs) + _, combined_kwargs = cls.split_url_kwargs(url) + + for k, v in combined_kwargs.items(): + if isinstance(v, str) and cls.PLUGIN_ATTR_PREFIX in v: + raise ValueError( + f"The URL parameter {k} depends on the plugin. " + "You must specify the value for this parameter " + "for this URL to be evaluated correctly. " + f"Try passing {k} as a keyword argument. " + f"e.g.: `URLConfig.evaluate_dry({url}, {k}=SOME_VALUE)`." + ) + + return cls.eval(url) diff --git a/strax/utils.py b/strax/utils.py index f080b460..0fa49dd9 100644 --- a/strax/utils.py +++ b/strax/utils.py @@ -3,9 +3,11 @@ from concurrent.futures import ThreadPoolExecutor, wait, FIRST_COMPLETED import itertools import contextlib +import inspect from functools import wraps import json import re +import reprlib import sys import traceback import typing as ty @@ -867,3 +869,90 @@ def convert_structured_array_to_df(structured_array, log=None): ) return pd.DataFrame(data_dict) + + +@export +def filter_kwargs(func, kwargs): + """Filter out keyword arguments that are not in the call signature of func and return filtered + kwargs dictionary.""" + params = inspect.signature(func).parameters + if any([str(p).startswith("**") for p in params.values()]): + # if func accepts wildcard kwargs, return all + return kwargs + return {k: v for k, v in kwargs.items() if k in params} + + +@export +class CacheDict(collections.OrderedDict): + """Dict with a limited length, ejecting LRUs as needed. + + copied from + https://gist.github.com/davesteele/44793cd0348f59f8fadd49d7799bd306 + + """ + + def __init__(self, *args, cache_len: int = 10, **kwargs): + assert cache_len > 0 + self.cache_len = cache_len + + super().__init__(*args, **kwargs) + + def __setitem__(self, key, value): + super().__setitem__(key, value) + super().move_to_end(key) + + while len(self) > self.cache_len: + oldkey = next(iter(self)) + super().__delitem__(oldkey) + + def __getitem__(self, key): + val = super().__getitem__(key) + super().move_to_end(key) + + return val + + +@export +def total_size(o, handlers=None, verbose=False): + """Returns the approximate memory footprint an object and all of its contents. + + Automatically finds the contents of the following builtin containers and + their subclasses: tuple, list, deque, dict, set and frozenset. + To search other containers, add handlers to iterate over their contents: + + handlers = {SomeContainerClass: iter, + OtherContainerClass: OtherContainerClass.get_elements} + + from: https://code.activestate.com/recipes/577504/ + + """ + dict_handler = lambda d: itertools.chain.from_iterable(d.items()) + all_handlers = { + tuple: iter, + list: iter, + collections.deque: iter, + dict: dict_handler, + set: iter, + frozenset: iter, + } + if handlers is not None: + all_handlers.update(handlers) # user handlers take precedence + seen = set() # track which object id's have already been seen + default_size = sys.getsizeof(0) # estimate sizeof object without __sizeof__ + + def sizeof(o): + if id(o) in seen: # do not double count the same object + return 0 + seen.add(id(o)) + s = sys.getsizeof(o, default_size) + + if verbose: + print(s, type(o), reprlib.repr(o), file=sys.stderr) + + for typ, handler in all_handlers.items(): + if isinstance(o, typ): + s += sum(map(sizeof, handler(o))) + break + return s + + return sizeof(o) diff --git a/tests/test_url_config.py b/tests/test_url_config.py new file mode 100644 index 00000000..4c29e445 --- /dev/null +++ b/tests/test_url_config.py @@ -0,0 +1,451 @@ +"""Tests for the URLConfig engine in strax.url_config. + +strax itself does not ship any protocols, so the tests register their own on a subclass with +separate registries (``_LOOKUP`` / ``_PREPROCESSORS``). This keeps the test protocols out of the +global ``strax.URLConfig`` registry that downstream packages (e.g. straxen) register into. + +""" + +import json +import pickle +import random +import unittest + +import numpy as np +import pandas as pd + +import strax + + +class DummyObject: + def __init__(self, **kwargs): + self.__dict__.update(kwargs) + + +class URLConfigForTests(strax.URLConfig): + """URLConfig with its own protocol and preprocessor registries.""" + + _LOOKUP: dict = {} + _PREPROCESSORS = () + + +@URLConfigForTests.register("json") +def read_json(content: str, **kwargs): + return json.loads(content) + + +@URLConfigForTests.register("format") +def format_arg(arg: str, **kwargs): + return arg.format(**kwargs) + + +@URLConfigForTests.register("take") +def get_key(container, take=None, **kwargs): + if take is None: + return container + if not isinstance(take, list): + take = [take] + for t in take: + container = container[t] + return container + + +@URLConfigForTests.register("random") +def generate_random(_): + return random.random() + + +@URLConfigForTests.register("unpicklable") +def return_lambda(_): + return lambda x: x + + +@URLConfigForTests.register("large-array") +def large_array(_): + return np.ones(1_000_000).tolist() + + +@URLConfigForTests.register("only-a") +def only_a(arg, a=None): + return (arg, a) + + +@URLConfigForTests.preprocessor +def formatter(config, **kwargs): + if not isinstance(config, str): + return config + try: + config = config.format(**kwargs) + except KeyError: + pass + return config + + +GLOBAL_VERSIONS = {"global_v1": {"test_config": "v0"}} + + +@URLConfigForTests.preprocessor +def replace_global_version(config, name=None, **kwargs): + if name is None: + return + + if not isinstance(config, str): + return config + + if URLConfigForTests.SCHEME_SEP not in config: + return config + + version = URLConfigForTests.kwarg_from_url(config, "version") + + if version is None: + return config + + if version.startswith("global_") and version in GLOBAL_VERSIONS: + version = GLOBAL_VERSIONS[version].get(name, version) + config = URLConfigForTests.format_url_kwargs(config, version=version) + return config + + +class ExamplePlugin(strax.Plugin): + depends_on = () + dtype = strax.time_fields + provides = "test_data" + test_config = URLConfigForTests(default=42) + cached_config = URLConfigForTests(default=666, cache=1) + typed_config = URLConfigForTests(default=1, type=int) + + +RUN_ID = "000001" + + +class TestURLConfig(unittest.TestCase): + def setUp(self): + self.st = strax.Context(storage=[], register=[ExamplePlugin]) + + def get_plugin(self): + return self.st.get_single_plugin(RUN_ID, "test_data") + + # --- Exports and registries --- + + def test_exported(self): + assert issubclass(strax.URLConfig, strax.Config) + assert callable(strax.clear_config_caches) + assert callable(strax.config_cache_size_mb) + + def test_registry_isolated(self): + """Test protocols must not leak into the global strax.URLConfig registry.""" + assert "random" in URLConfigForTests._LOOKUP + assert "random" not in strax.URLConfig._LOOKUP + assert formatter not in [f for _, f in strax.URLConfig._PREPROCESSORS] + + def test_register_duplicate(self): + with self.assertRaises(ValueError): + URLConfigForTests.register("json", read_json) + + def test_register_non_string(self): + with self.assertRaises(ValueError): + URLConfigForTests.register(42, read_json) + + def test_register_tuple(self): + class LocalURLConfig(strax.URLConfig): + _LOOKUP: dict = {} + _PREPROCESSORS = () + + @LocalURLConfig.register(("proto-a", "proto-b")) + def ident(arg): + return arg + + assert LocalURLConfig._LOOKUP["proto-a"] is ident + assert LocalURLConfig._LOOKUP["proto-b"] is ident + + def test_preprocessor_duplicate(self): + with self.assertRaises(ValueError): + URLConfigForTests.preprocessor(formatter) + + def test_preprocessor_precedence(self): + """Preprocessors with higher precedence run first.""" + calls = [] + + class LocalURLConfig(strax.URLConfig): + _LOOKUP: dict = {} + _PREPROCESSORS = () + + @LocalURLConfig.preprocessor(precedence=0) + def low(config): + calls.append("low") + return config + "-low" + + @LocalURLConfig.preprocessor(precedence=10) + def high(config): + calls.append("high") + return config + "-high" + + class LocalPlugin(ExamplePlugin): + local_config = LocalURLConfig(default="x") + + st = strax.Context(storage=[], register=[LocalPlugin]) + p = st.get_single_plugin(RUN_ID, "test_data") + # validate may run more than once per plugin; each pass runs high before low + self.assertEqual(calls[:2], ["high", "low"]) + self.assertEqual(calls, ["high", "low"] * (len(calls) // 2)) + self.assertEqual(p.local_config, "x-high-low") + + # --- Fetching values through a plugin --- + + def test_default(self): + p = self.get_plugin() + self.assertEqual(p.test_config, 42) + + def test_literal(self): + self.st.set_config({"test_config": 666}) + p = self.get_plugin() + self.assertEqual(p.test_config, 666) + + def test_unknown_protocol_is_literal(self): + self.st.set_config({"test_config": "unknown://abc"}) + p = self.get_plugin() + self.assertEqual(p.test_config, "unknown://abc") + + def test_leading_zero_int(self): + self.st.set_config({"test_config": "format://{value}?value=0666"}) + p = self.get_plugin() + self.assertEqual(p.test_config, "0666") + self.assertIsInstance(p.test_config, str) + + def test_json_protocol(self): + self.st.set_config({"test_config": 'json://{"a":0}'}) + p = self.get_plugin() + self.assertEqual(p.test_config, {"a": 0}) + + def test_plugin_attribute_lookup(self): + self.st.set_config({"test_config": "format://{run_id}?run_id=plugin.run_id"}) + p = self.get_plugin() + self.assertEqual(p.test_config, RUN_ID) + + def test_config_lookup(self): + self.st.set_config({"test_config": "format://{v}?v=config.cached_config"}) + p = self.get_plugin() + self.assertEqual(p.test_config, "666") + + def test_take_nested(self): + self.st.set_config( + { + "test_config": ( + 'take://json://{"a":{"aa":0,"ab":1},"b":{"ba":2,"bb":3}}?take=b&take=ba' + ) + } + ) + p = self.get_plugin() + self.assertEqual(p.test_config, 2) + + def test_eval_filters_kwargs(self): + """Kwargs that the protocol does not accept are dropped.""" + self.assertEqual(URLConfigForTests.eval("only-a://x?a=1&b=2"), ("x", 1)) + + def test_format_preprocessor(self): + self.st.set_config({"test_config": "{name}:{run_id}"}) + p = self.get_plugin() + self.assertEqual(p.test_config, f"test_config:{RUN_ID}") + self.assertEqual(p.test_config, p.config["test_config"]) + + def test_global_version_preprocessor(self): + self.st.set_config({"test_config": "fake://url?version=global_v1"}) + p = self.get_plugin() + self.assertEqual(p.test_config, "fake://url?version=v0") + + def test_global_version_not_changed(self): + """ + - if no global version is matched, the url version should not be changed + - if config is not matched, the url version should not be changed + """ + assert "global_v2" not in GLOBAL_VERSIONS + self.st.set_config({"test_config": "fake://url?version=global_v2"}) + p = self.get_plugin() + self.assertEqual(p.test_config, "fake://url?version=global_v2") + + class ExamplePluginNew(ExamplePlugin): + test_config_new = URLConfigForTests(default="rootisthesourceofallevil") + + self.st.register(ExamplePluginNew) + self.st.set_config({"test_config_new": "fake://url?version=global_v1"}) + p = self.get_plugin() + self.assertEqual(p.test_config_new, "fake://url?version=global_v1") + + def test_url_is_hashed_not_value(self): + """The lineage contains the (preprocessed) URL, not the fetched value.""" + url = 'take://json://{"a":1}?take=a' + self.st.set_config({"test_config": url}) + lineage = self.st.key_for(RUN_ID, "test_data").lineage + self.assertEqual(lineage["test_data"][2]["test_config"], url) + + # --- Type validation --- + + def test_type_not_enforced_on_url(self): + self.st.set_config({"typed_config": 'json://"2"'}) + p = self.get_plugin() + # pydantic coerces "2" -> 2 is allowed, the value is returned as-is + self.assertEqual(p.typed_config, "2") + + def test_type_validated_on_fetched_value(self): + self.st.set_config({"typed_config": 'json://"abc"'}) + p = self.get_plugin() + with self.assertRaises(TypeError): + p.typed_config + + def test_type_validated_on_literal(self): + self.st.set_config({"typed_config": "abc"}) + with self.assertRaises(TypeError): + self.get_plugin() + + def test_type_tuple(self): + class LocalPlugin(ExamplePlugin): + tuple_typed = URLConfigForTests(default=1, type=(int, dict)) + + st = strax.Context(storage=[], register=[LocalPlugin]) + st.set_config({"tuple_typed": 'json://{"a":1}'}) + self.assertEqual(st.get_single_plugin(RUN_ID, "test_data").tuple_typed, {"a": 1}) + st.set_config({"tuple_typed": 'json://"abc"'}) + with self.assertRaises(TypeError): + st.get_single_plugin(RUN_ID, "test_data").tuple_typed + + # --- Caching --- + + def test_cache(self): + p = self.get_plugin() + + # sanity check that default value is not affected + self.assertEqual(p.cached_config, 666) + self.st.set_config({"cached_config": "random://abc"}) + p = self.get_plugin() + + # value is randomly generated when accessed so if + # its equal when we access it again, its coming from the cache + cached_value = p.cached_config + self.assertEqual(cached_value, p.cached_config) + + # now change the config to which will generate a new number + self.st.set_config({"cached_config": "random://dfg"}) + p = self.get_plugin() + + # sanity check that the new value is still consistent i.e. cached + self.assertEqual(p.cached_config, p.cached_config) + + # test if previous value is evicted, since cache size is 1 + self.assertNotEqual(cached_value, p.cached_config) + + # verify pickalibility of objects in cache dont affect plugin pickalibility + self.st.set_config({"cached_config": "unpicklable://dfg"}) + p = self.get_plugin() + with self.assertRaises(AttributeError): + pickle.dumps(p.cached_config) + pickle.dumps(p) + + def test_uncached(self): + self.st.set_config({"test_config": "random://abc"}) + p = self.get_plugin() + self.assertNotEqual(p.test_config, p.test_config) + + def test_cache_size(self): + """Test the cache helper functions.""" + # make sure the value has a detectable size + self.st.set_config({"cached_config": "large-array://dfg"}) + p = self.get_plugin() + + # fetch the value so its stored in the cache + p.cached_config + + # cache should now have finite size + self.assertGreater(strax.config_cache_size_mb(), 0.0) + + # test if clearing cache works as expected + strax.clear_config_caches() + self.assertEqual(strax.config_cache_size_mb(), 0.0) + + # --- URL <-> AST conversion --- + + def test_url_to_ast(self): + ast = URLConfigForTests.url_to_ast('take://json://{"a":1}?take=a') + # kwargs are passed down to nested protocols that accept them (json takes **kwargs) + self.assertEqual(ast, ("take", ("json", '{"a":1}', {"take": "a"}), {"take": "a"})) + # ... and are dropped for protocols that do not accept them + ast = URLConfigForTests.url_to_ast("take://only-a://x?a=1&take=b") + # (take accepts **kwargs and keeps both) + self.assertEqual(ast, ("take", ("only-a", "x", {"a": 1}), {"a": 1, "take": "b"})) + + def test_url_to_ast_unknown_protocol(self): + self.assertEqual(URLConfigForTests.url_to_ast("unknown://abc"), (None, "unknown://abc", {})) + self.assertEqual(URLConfigForTests.url_to_ast("abc"), (None, "abc", {})) + with self.assertRaises(TypeError): + URLConfigForTests.url_to_ast(42) + + def test_ast_to_url(self): + # non-string arguments are wrapped in json:// + url = URLConfigForTests.ast_to_url("take", {"a": 1}, {"take": "a"}) + self.assertEqual(url, 'take://json://{"a": 1}?take=a') + self.assertEqual(URLConfigForTests.ast_to_url("take", 1.5), "take://json://1.5") + # nested AST tuples are supported + url = URLConfigForTests.ast_to_url(("take", ("format", "{x}", {"x": 1}), {})) + self.assertEqual(url, "take://format://{x}?x=1") + # protocol None returns the argument as-is + self.assertEqual(URLConfigForTests.ast_to_url(None, "abc"), "abc") + with self.assertRaises(TypeError): + URLConfigForTests.ast_to_url("take", object()) + + def test_ast_roundtrip(self): + url = 'take://json://{"a": {"b": 1}}?take=a&take=b' + ast = URLConfigForTests.url_to_ast(url) + self.assertEqual(URLConfigForTests.ast_to_url(*ast), url) + self.assertEqual(URLConfigForTests.eval(url), 1) + + def test_are_equal(self): + self.assertTrue( + URLConfigForTests.are_equal( + "format://{a}{b}?a=1&b=2", + "format://{a}{b}?b=2&a=1", + ) + ) + self.assertFalse(URLConfigForTests.are_equal("format://x?a=1", "format://x?a=2")) + + def test_url_kwargs(self): + url = "format://x?b=2&a=1&c=1&c=2&d=text" + path, kwargs = URLConfigForTests.split_url_kwargs(url) + self.assertEqual(path, "format://x") + self.assertEqual(kwargs, {"b": 2, "a": 1, "c": [1, 2], "d": "text"}) + self.assertEqual(URLConfigForTests.kwarg_from_url(url, "c"), [1, 2]) + self.assertIsNone(URLConfigForTests.kwarg_from_url(url, "missing")) + # kwargs are sorted, new kwargs override existing ones + self.assertEqual( + URLConfigForTests.format_url_kwargs(url, d="new"), + "format://x?a=1&b=2&c=1&c=2&d=new", + ) + + # --- Dry evaluation --- + + def test_dry_evaluation(self): + """A dry evaluation outside of a plugin yields the same result as in the plugin.""" + plugin_url = "format://{run_id}?run_id=plugin.run_id" + self.st.set_config({"test_config": plugin_url}) + correct_val = self.get_plugin().test_config + + dry_val1 = URLConfigForTests.evaluate_dry(f"format://{{run_id}}?run_id={RUN_ID}") + dry_val2 = URLConfigForTests.evaluate_dry("format://{run_id}", run_id=RUN_ID) + self.assertEqual(correct_val, dry_val1) + self.assertEqual(correct_val, dry_val2) + + # dry-evaluation can not resolve plugin attributes + with self.assertRaises(ValueError): + URLConfigForTests.evaluate_dry(plugin_url) + + # --- Descriptions --- + + def test_protocol_descr(self): + df = URLConfigForTests.protocol_descr() + self.assertIsInstance(df, pd.DataFrame) + self.assertIn("json://", df["name"].values) + URLConfigForTests.print_protocols() + + def test_preprocessor_descr(self): + df = URLConfigForTests.preprocessor_descr() + self.assertEqual(len(df), len(URLConfigForTests._PREPROCESSORS)) + URLConfigForTests.print_preprocessors() + URLConfigForTests.print_summary() diff --git a/tests/test_utils.py b/tests/test_utils.py index 329ae0ec..9cf23112 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -1,6 +1,7 @@ import numpy as np import tempfile import unittest +import pytest import shutil import hypothesis.strategies as hst from hypothesis import given, settings @@ -242,3 +243,40 @@ def _test_get_array_multi_run(self, max_worker=None): def tearDown(self): shutil.rmtree(self.tempdir) + + +def test_filter_kwargs(): + all_kwargs = dict(a=1, b=2, c=3) + + def func1(a=None, b=None): + return + + def func2(a=None, **kwargs): + return + + assert strax.filter_kwargs(func1, all_kwargs) == dict(a=1, b=2) + # func2 accepts wildcard kwargs, so nothing is filtered + assert strax.filter_kwargs(func2, all_kwargs) == all_kwargs + + +def test_cache_dict(): + cache = strax.CacheDict(cache_len=2) + cache["a"] = 1 + cache["b"] = 2 + # accessing "a" makes "b" the least recently used entry + assert cache["a"] == 1 + cache["c"] = 3 + assert list(cache) == ["a", "c"] + assert "b" not in cache + + with pytest.raises(AssertionError): + strax.CacheDict(cache_len=0) + + +def test_total_size(): + small = strax.total_size([1]) + big = strax.total_size([list(range(1000))]) + assert 0 < small < big + # shared objects are counted only once + shared = list(range(1000)) + assert strax.total_size([shared, shared]) < 2 * strax.total_size([shared])