Copy over xarray's set_options context manager (#243)

Author: kthyng

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()

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]]
 

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)

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():

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)