diff --git a/cf_xarray/__init__.py b/cf_xarray/__init__.py index e559c639..f80585e6 100644 --- a/cf_xarray/__init__.py +++ b/cf_xarray/__init__.py @@ -1,5 +1,6 @@ from .accessor import CFAccessor # noqa from .helpers import bounds_to_vertices, vertices_to_bounds # noqa +from .options import set_options # noqa from .utils import _get_version __version__ = _get_version() diff --git a/cf_xarray/accessor.py b/cf_xarray/accessor.py index df85f5af..4b0f7ce7 100644 --- a/cf_xarray/accessor.py +++ b/cf_xarray/accessor.py @@ -27,6 +27,7 @@ from .criteria import cf_role_criteria, coordinate_criteria, regex from .helpers import bounds_to_vertices +from .options import OPTIONS from .utils import ( _get_version, _is_datetime_like, @@ -65,13 +66,6 @@ ATTRS["time"] = ATTRS["T"] ATTRS["vertical"] = ATTRS["Z"] -OPTIONS: MutableMapping[str, Any] = {"custom_criteria": []} - - -def set_options(custom_criteria): - OPTIONS["custom_criteria"] = always_iterable(custom_criteria, allowed=(tuple, list)) - - # Type for Mapper functions Mapper = Callable[[Union[DataArray, Dataset], str], List[str]] diff --git a/cf_xarray/options.py b/cf_xarray/options.py new file mode 100644 index 00000000..991d8be3 --- /dev/null +++ b/cf_xarray/options.py @@ -0,0 +1,55 @@ +""" +Started from xarray options.py +""" + +import copy +from typing import Any, MutableMapping + +from .utils import always_iterable + +OPTIONS: MutableMapping[str, Any] = { + "custom_criteria": [], +} + + +class set_options: + """Set options for cf-xarray in a controlled context. + Currently supported options: + - ``custom_critera``: Translate from axis, coord, or custom name to + variable name optionally using ``custom_criteria``. Default: []. + + You can use ``set_options`` either as a context manager: + >>> my_custom_criteria = { 'ssh': {'name': 'elev$'} } + >>> ds = xr.Dataset({"elev": np.arange(1000)}) + >>> with cf_xarray.set_options(custom_criteria=my_custom_criteria): + ... assert (ds['elev'] == ds.cf['ssh']).all() + + Or to set global options: + >>> cf_xarray.set_options(custom_criteria=my_custom_criteria) + >>> assert (ds['elev'] == ds.cf['ssh']).all() + """ + + def __init__(self, **kwargs): + self.old = {} + for k, v in kwargs.items(): + if k not in OPTIONS: + raise ValueError( + f"argument name {k!r} is not in the set of valid options {set(OPTIONS)!r}" + ) + self.old[k] = OPTIONS[k] + self._apply_update(kwargs) + + def _apply_update(self, options_dict): + options_dict = copy.deepcopy(options_dict) + for k, v in options_dict.items(): + if k == "custom_criteria": + options_dict["custom_criteria"] = always_iterable( + options_dict["custom_criteria"], allowed=(tuple, list) + ) + OPTIONS.update(options_dict) + + def __enter__(self): + return + + def __exit__(self, type, value, traceback): + self._apply_update(self.old) diff --git a/cf_xarray/tests/test_accessor.py b/cf_xarray/tests/test_accessor.py index c4a33abe..ecb4ea52 100644 --- a/cf_xarray/tests/test_accessor.py +++ b/cf_xarray/tests/test_accessor.py @@ -1252,10 +1252,11 @@ def test_custom_criteria(): }, } my_custom_criteria2 = {"temp": {"name": "temperature"}} - cf_xarray.accessor.set_options(my_custom_criteria) my_custom_criteria_list = [my_custom_criteria, my_custom_criteria2] my_custom_criteria_tuple = (my_custom_criteria, my_custom_criteria2) + cf_xarray.set_options(custom_criteria=my_custom_criteria) + # Match by name regex match ds = xr.Dataset() ds["salinity"] = ("dim", np.arange(10)) @@ -1301,16 +1302,16 @@ def test_custom_criteria(): ) # test criteria list of dicts - cf_xarray.accessor.set_options(my_custom_criteria_list) - ds = xr.Dataset() - ds["temperature"] = ("dim", np.arange(10)) - assert_identical(ds.cf["temp"], ds["temperature"]) + with cf_xarray.set_options(custom_criteria=my_custom_criteria_list): + ds = xr.Dataset() + ds["temperature"] = ("dim", np.arange(10)) + assert_identical(ds.cf["temp"], ds["temperature"]) # test criteria tuple of dicts - cf_xarray.accessor.set_options(my_custom_criteria_tuple) - ds = xr.Dataset() - ds["temperature"] = ("dim", np.arange(10)) - assert_identical(ds.cf["temp"], ds["temperature"]) + with cf_xarray.set_options(custom_criteria=my_custom_criteria_tuple): + ds = xr.Dataset() + ds["temperature"] = ("dim", np.arange(10)) + assert_identical(ds.cf["temp"], ds["temperature"]) def test_cf_standard_name_table_version(): diff --git a/cf_xarray/tests/test_options.py b/cf_xarray/tests/test_options.py new file mode 100644 index 00000000..786ac337 --- /dev/null +++ b/cf_xarray/tests/test_options.py @@ -0,0 +1,14 @@ +""" +Tests OPTIONS logic brought in from xarray. +""" + +import pytest + +import cf_xarray as cfxr + + +def test_options(): + + # test for inputting a nonexistent option + with pytest.raises(ValueError): + cfxr.set_options(DISPLAY_WIDTH=80)