scikit-hep · matthewfeickert · Mar 23, 2022 · Mar 23, 2022 · Mar 23, 2022 · Mar 23, 2022
diff --git a/src/pyhf/schema/__init__.py b/src/pyhf/schema/__init__.py
@@ -23,17 +23,40 @@ class Schema(sys.modules[__name__].__class__):
     """
     A module-level wrapper around :mod:`pyhf.schema` which will provide additional functionality for interacting with schemas.
 
-    Example:
+    .. rubric:: Example (callable)
+
+    .. code-block:: pycon
+
         >>> import pyhf.schema
         >>> import pathlib
+        >>> new_path = pathlib.Path("/home/root/my/new/path")
         >>> curr_path = pyhf.schema.path
-        >>> curr_path # doctest: +ELLIPSIS
+        >>> curr_path  # doctest: +ELLIPSIS
         PosixPath('.../pyhf/schemas')
-        >>> pyhf.schema(pathlib.Path('/home/root/my/new/path'))
+        >>> pyhf.schema(new_path)  # doctest: +ELLIPSIS
+        <module 'pyhf.schema' from ...>
         >>> pyhf.schema.path
         PosixPath('/home/root/my/new/path')
-        >>> pyhf.schema(curr_path)
-        >>> pyhf.schema.path # doctest: +ELLIPSIS
+        >>> pyhf.schema(curr_path)  # doctest: +ELLIPSIS
+        <module 'pyhf.schema' from ...>
+        >>> pyhf.schema.path  # doctest: +ELLIPSIS
+        PosixPath('.../pyhf/schemas')
+
+    .. rubric:: Example (context-manager)
+
+    .. code-block:: pycon
+
+        >>> import pyhf.schema
+        >>> import pathlib
+        >>> new_path = pathlib.Path("/home/root/my/new/path")
+        >>> curr_path = pyhf.schema.path
+        >>> curr_path  # doctest: +ELLIPSIS
+        PosixPath('.../pyhf/schemas')
+        >>> with pyhf.schema(new_path):
+        ...     print(repr(pyhf.schema.path))
+        ...
+        PosixPath('/home/root/my/new/path')
+        >>> pyhf.schema.path  # doctest: +ELLIPSIS
         PosixPath('.../pyhf/schemas')
 
     """
@@ -45,10 +68,23 @@ def __call__(self, new_path: pathlib.Path):
         Args:
             new_path (pathlib.Path): Path to folder containing the schemas
 
+        Returns:
+            self (pyhf.schema.Schema): Returns itself (for contextlib management)
+        """
+        self.orig_path, variables.schemas = variables.schemas, new_path
+        return self
+
+    def __enter__(self):
+        pass
+
+    def __exit__(self, *args, **kwargs):
+        """
+        Reset the local search path for finding schemas locally.
+
         Returns:
             None
         """
-        variables.schemas = new_path
+        variables.schemas = self.orig_path
 
     @property
     def path(self):

diff --git a/tests/test_public_api.py b/tests/test_public_api.py
@@ -1,6 +1,7 @@
 import pytest
 import pyhf
 import numpy as np
+import pathlib
 
 
 @pytest.fixture(scope='function')
@@ -200,3 +201,23 @@ def test_pdf_batched(backend):
 
     model.pdf(pars, data)
     model.expected_data(pars)
+
+
+def test_set_schema_path(monkeypatch):
+    monkeypatch.setattr(
+        pyhf.schema.variables, 'schemas', pyhf.schema.variables.schemas, raising=True
+    )
+
+    new_path = pathlib.Path('a/new/path')
+    pyhf.schema(new_path)
+    assert pyhf.schema.path == new_path
+
+
+def test_set_schema_path_context(monkeypatch):
+    monkeypatch.setattr(
+        pyhf.schema.variables, 'schemas', pyhf.schema.variables.schemas, raising=True
+    )
+
+    new_path = pathlib.Path('a/new/path')
+    with pyhf.schema(new_path):
+        assert pyhf.schema.path == new_path
diff --git a/tests/test_schema.py b/tests/test_schema.py
@@ -27,18 +27,51 @@ def test_schema_callable():
     assert callable(pyhf.schema)
 
 
-def test_schema_changeable(datadir):
+def test_schema_changeable(datadir, monkeypatch):
+    monkeypatch.setattr(
+        pyhf.schema.variables, 'schemas', pyhf.schema.variables.schemas, raising=True
+    )
+    old_path = pyhf.schema.path
+    new_path = datadir / 'customschema'
+
     with pytest.raises(pyhf.exceptions.SchemaNotFound):
         pyhf.Workspace(json.load(open(datadir / 'customschema' / 'custom.json')))
 
-    old_path = pyhf.schema.path
-    pyhf.schema(datadir / 'customschema')
-    assert pyhf.schema.path != old_path
-    assert pyhf.schema.path == datadir / 'customschema'
-    assert pyhf.Workspace(json.load(open(datadir / 'customschema' / 'custom.json')))
+    pyhf.schema(new_path)
+    assert old_path != pyhf.schema.path
+    assert new_path == pyhf.schema.path
+    assert pyhf.Workspace(json.load(open(new_path / 'custom.json')))
     pyhf.schema(old_path)
 
 
+def test_schema_changeable_context(datadir, monkeypatch):
+    monkeypatch.setattr(
+        pyhf.schema.variables, 'schemas', pyhf.schema.variables.schemas, raising=True
+    )
+    old_path = pyhf.schema.path
+    new_path = datadir / 'customschema'
+
+    assert old_path == pyhf.schema.path
+    with pyhf.schema(new_path):
+        assert old_path != pyhf.schema.path
+        assert new_path == pyhf.schema.path
+        assert pyhf.Workspace(json.load(open(new_path / 'custom.json')))
+    assert old_path == pyhf.schema.path
+
+
+def test_schema_changeable_context_error(datadir, monkeypatch):
+    monkeypatch.setattr(
+        pyhf.schema.variables, 'schemas', pyhf.schema.variables.schemas, raising=True
+    )
+    old_path = pyhf.schema.path
+    new_path = datadir / 'customschema'
+
+    with pytest.raises(ZeroDivisionError):
+        with pyhf.schema(new_path):
+            raise ZeroDivisionError()
+    assert old_path == pyhf.schema.path
+
+
 def test_no_channels():
     spec = {'channels': []}
     with pytest.raises(pyhf.exceptions.InvalidSpecification):