Repository: scikit-hep/pyhf

Issue:
Add contextlib support to the pyhf.schema API
> instead of having users overwriting the value by assignment, do you think there is value in offering a `pyhf.utils.update_schema_path` or something that performs this operation through the API?
>
> Thanks @kratsg for taking up my suggestion!
> 
> I second @matthewfeickert's suggestion to change schemas via function call rather than assignment. It is much simpler to add functionality to a function if it ever becomes necessary, than to replace an entire existing module with a class instance with a property.
> 
> I'd even go so far to say that a context manager doubling as an update function would be ideal IMO:
> ```python
> # in pyhf.utils
> _SCHEMAS = Path(...)
> class use_schema_path: # snake_case to remind of function-like usage
>     def  __init__(self, path):
>         global _SCHEMAS
>         self._old_schemas = _SCHEMAS
>         _SCHEMAS = pathlib.Path(path)
>     def __enter__(self):
>         pass
>     def __exit__(self, *args, **kwargs):
>         global _SCHEMAS
>         _SCHEMAS = self._old_schemas
> ```
> which can still be called as a function (only executing `__init__`), so short scripts etc. are not forced to use `with` blocks.
> But it can also be used like so:
> ```python
> def make_my_workspace(spec):
>     with pyhf.utils.use_schema_path('/my/very/special/schemas'):
>         return pyhf.Workspace(spec)
> ```
> So as a user writing code on top of pyhf, I don't have to worry about resesetting the global variable, the CM does it for me, and there are fewer mistakes to make.
> 
_Originally posted by @lhenkelm in https://github.com/scikit-hep/pyhf/issues/1753#issuecomment-1026678066_


Discussion:
This is on top of #1609 .

Resolving patch:
diff --git a/src/pyhf/schema/__init__.py b/src/pyhf/schema/__init__.py
--- 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
         >>> 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'))
+        >>> new_path = 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
+        >>> curr_path = pyhf.schema.path
+        >>> curr_path  # doctest: +ELLIPSIS
+        PosixPath('.../pyhf/schemas')
+        >>> new_path = pathlib.Path("/home/root/my/new/path")
+        >>> 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):


Test patch:
diff --git a/tests/test_public_api.py b/tests/test_public_api.py
--- 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
--- 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):