Repository: gaogaotiantian/coredumpy

Issue:
Support a similar switch to pytest as unittest
The users should be able to do `coredumpy.patch_pytest` as well.

Hints from issue comments:

Solution patch:
diff --git a/src/coredumpy/__init__.py b/src/coredumpy/__init__.py
index db4fc09..6127d27 100644
--- a/src/coredumpy/__init__.py
+++ b/src/coredumpy/__init__.py
@@ -9,6 +9,7 @@ from .config import config
 from .coredumpy import Coredumpy, dump, dumps, load
 from .except_hook import patch_except
 from .main import main
+from .pytest_hook import patch_pytest
 from .type_support import TypeSupportBase, TypeSupportContainerBase, NotReady
 from .unittest_hook import patch_unittest
 from .conf_hook import startup_conf
@@ -24,6 +25,7 @@ __all__ = [
     "load",
     "main",
     "patch_except",
+    "patch_pytest",
     "patch_unittest",
     "pytest_hook",
     "TypeSupportBase",


Test patch:
diff --git a/src/coredumpy/pytest_hook.py b/src/coredumpy/pytest_hook.py
index 5c5d08c..ce5e12a 100644
--- a/src/coredumpy/pytest_hook.py
+++ b/src/coredumpy/pytest_hook.py
@@ -1,6 +1,23 @@
+# Licensed under the Apache License: http://www.apache.org/licenses/LICENSE-2.0
+# For details: https://github.com/gaogaotiantian/coredumpy/blob/master/NOTICE.txt
+
+from typing import Optional
+
 import coredumpy
 
 
+_enable_coredumpy = False
+_coredumpy_dir = "./coredumpy"
+
+
+def patch_pytest(directory: Optional[str] = None):
+    global _enable_coredumpy
+    global _coredumpy_dir
+    _enable_coredumpy = True
+    if directory is not None:
+        _coredumpy_dir = directory
+
+
 def pytest_addoption(parser):
     parser.addoption(
         "--enable-coredumpy",
@@ -23,7 +40,11 @@ def _get_description(report):
 
 
 def pytest_exception_interact(node, call, report):
-    if not node.config.getoption("--enable-coredumpy"):
+    if node.config.getoption("--enable-coredumpy"):
+        directory = node.config.getoption("--coredumpy-dir")
+    elif _enable_coredumpy:
+        directory = _coredumpy_dir
+    else:
         return
 
     import pytest  # type: ignore
@@ -34,7 +55,7 @@ def pytest_exception_interact(node, call, report):
                 tb = tb.tb_next
             filename = coredumpy.dump(tb.tb_frame,
                                       description=_get_description(report),
-                                      directory=node.config.getoption("--coredumpy-dir"))
+                                      directory=directory)
             print(f'Your frame stack is dumped, open it with\n'
                   f'coredumpy load {filename}')
         except Exception:  # pragma: no cover
diff --git a/tests/test_pytest.py b/tests/test_pytest.py
index 62acad6..5bfd60c 100644
--- a/tests/test_pytest.py
+++ b/tests/test_pytest.py
@@ -43,3 +43,30 @@ class TestPytest(TestBase):
             """
             self.run_script(script)
             self.assertEqual(len(os.listdir(dump_path)), 2)
+
+    def test_pytest_patch(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            test = textwrap.dedent("""
+                def test_for_pytest_equal():
+                    assert 1 == 2
+
+                def test_for_pytest_greater():
+                    assert 1 > 2
+            """)
+            with open(os.path.join(tmpdir, "test.py"), "w") as f:
+                f.write(test)
+
+            test_path = os.path.join(tmpdir, "test.py")
+            dump_path = os.path.join(tmpdir, "dump")
+
+            script = f"""
+                import coredumpy
+                import pytest
+                import os
+                rootdir = os.path.dirname(__file__)
+                coredumpy.patch_pytest(directory={repr(dump_path)})
+                pytest.main(["--rootdir", rootdir, {repr(test_path)}])
+            """
+            stdout, stderr = self.run_script(script)
+            self.assertEqual(len(os.listdir(dump_path)), 2,
+                             f"The dump directory has {os.listdir(dump_path)}\n{stdout}\n{stderr}")