add pytest_raises_on_rank (#162)

* add pytest_raises_on_rank

* better type annotation

* add copyright to mpi.py

* add test_mpi.py

* fix doc

* Clarify test

doc/conf.py

@@ -29,8 +29,11 @@ intersphinx_mapping = {
     "https://numpy.org/doc/stable": None,
     "https://documen.tician.de/pymbolic/": None,
     "https://documen.tician.de/loopy/": None,
+    "https://docs.pytest.org/en/stable/": None,
     }
 
 nitpick_ignore_regex = [
         ["py:class", r"typing_extensions\.(.+)"],
         ]
+
+nitpicky = True

doc/index.rst

@@ -11,6 +11,7 @@ Welcome to pytools's documentation!
     graph
     tag
     codegen
+    mpi
     misc
     🚀 Github <https://github.com/inducer/pytools>
     💾 Download Releases <https://pypi.python.org/pypi/pytools>

doc/mpi.rst

+.. automodule:: pytools.mpi

pytools/mpi.py

+__copyright__ = """
+Copyright (C) 2009-2019 Andreas Kloeckner
+Copyright (C) 2022 University of Illinois Board of Trustees
+"""
+
+__license__ = """
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.
+"""
+
+__doc__ = """
+MPI helper functionality
+========================
+
+.. autofunction:: check_for_mpi_relaunch
+.. autofunction:: run_with_mpi_ranks
+.. autofunction:: pytest_raises_on_rank
+"""
+
+from contextlib import contextmanager, AbstractContextManager
+from typing import Generator, Tuple, Union, Type
+
+
 def check_for_mpi_relaunch(argv):
     if argv[1] != "--mpi-relaunch":
         return
@@ -26,3 +64,23 @@ def run_with_mpi_ranks(py_script, ranks, callable_, args=(), kwargs=None):
     check_call(["mpirun", "-np", str(ranks),
         sys.executable, py_script, "--mpi-relaunch", callable_and_args],
         env=newenv)
+
+
+@contextmanager
+def pytest_raises_on_rank(my_rank: int, fail_rank: int,
+        expected_exception: Union[Type[BaseException],
+                                  Tuple[Type[BaseException], ...]]) \
+                -> Generator[AbstractContextManager, None, None]:
+    """
+    Like :func:`pytest.raises`, but only expect an exception on rank *fail_rank*.
+    """
+    import pytest
+    from contextlib import nullcontext
+
+    if my_rank == fail_rank:
+        cm: AbstractContextManager = pytest.raises(expected_exception)
+    else:
+        cm = nullcontext()
+
+    with cm as exc:
+        yield exc

test/test_mpi.py

+__copyright__ = "Copyright (C) 2022 University of Illinois Board of Trustees"
+
+__license__ = """
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.
+"""
+
+import pytest
+
+
+def test_pytest_raises_on_rank():
+    from pytools.mpi import pytest_raises_on_rank
+
+    def fail(my_rank: int, fail_rank: int) -> None:
+        if my_rank == fail_rank:
+            raise ValueError("test failure")
+
+    with pytest.raises(ValueError):
+        fail(0, 0)
+
+    fail(0, 1)
+
+    with pytest_raises_on_rank(0, 0, ValueError):
+        # Generates an exception, and pytest_raises_on_rank
+        # expects one.
+        fail(0, 0)
+
+    with pytest_raises_on_rank(0, 1, ValueError):
+        # Generates no exception, and pytest_raises_on_rank
+        # does not expect one.
+        fail(0, 1)