factor out _set_hooks

src/_pytest/assertion/__init__.py

@@ -12,9 +12,13 @@
 from _pytest.compat import TYPE_CHECKING
 from _pytest.config import Config
 from _pytest.config import hookimpl
+from _pytest.monkeypatch import MonkeyPatch
 
 if TYPE_CHECKING:
+    from typing import Generator
+
     from _pytest.main import Session
+    from _pytest.nodes import Item
 
 
 def pytest_addoption(parser):
@@ -110,20 +114,12 @@ def pytest_collection(session: "Session") -> None:
             assertstate.hook.set_session(session)
 
 
-@hookimpl(tryfirst=True, hookwrapper=True)
-def pytest_runtest_protocol(item):
-    """Setup the pytest_assertrepr_compare and pytest_assertion_pass hooks
-
-    The newinterpret and rewrite modules will use util._reprcompare if
-    it exists to use custom reporting via the
-    pytest_assertrepr_compare hook.  This sets up this custom
-    comparison for the test.
-    """
-
+def _set_hooks(item: "Item") -> "MonkeyPatch":
     ihook = item.ihook
 
-    def callbinrepr(op, left, right):
-        # type: (str, object, object) -> Optional[str]
+    mp = MonkeyPatch()
+
+    def callbinrepr(op: str, left: object, right: object) -> Optional[str]:
         """Call the pytest_assertrepr_compare hook and prepare the result
 
         This uses the first result from the hook and then ensures the
@@ -148,19 +144,30 @@ def callbinrepr(op, left, right):
             return res
         return None
 
-    saved_assert_hooks = util._reprcompare, util._assertion_pass
-    util._reprcompare = callbinrepr
+    mp.setattr(util, "_reprcompare", callbinrepr)
 
     if ihook.pytest_assertion_pass.get_hookimpls():
 
         def call_assertion_pass_hook(lineno, orig, expl):
             ihook.pytest_assertion_pass(item=item, lineno=lineno, orig=orig, expl=expl)
 
-        util._assertion_pass = call_assertion_pass_hook
+        mp.setattr(util, "_assertion_pass", call_assertion_pass_hook)
 
-    yield
+    return mp
+
+
+@hookimpl(tryfirst=True, hookwrapper=True)
+def pytest_runtest_protocol(item: "Item") -> "Generator[None, None, None]":
+    """Setup the pytest_assertrepr_compare and pytest_assertion_pass hooks
 
-    util._reprcompare, util._assertion_pass = saved_assert_hooks
+    The newinterpret and rewrite modules will use util._reprcompare if
+    it exists to use custom reporting via the
+    pytest_assertrepr_compare hook.  This sets up this custom
+    comparison for the test.
+    """
+    mp = _set_hooks(item)
+    yield
+    mp.undo()
 
 
 def pytest_sessionfinish(session):