Add unittest support and fix directory (#11)

README.md

@@ -15,7 +15,17 @@ coredumpy saves your crash site so you can better debug your python program.
 
 ### dump
 
-You can dump the current frame stack by
+In most cases, you only need to hook `coredumpy` to some triggers
+
+```python
+import coredumpy
+# Create a dump in "./dumps" when there's an unhandled exception
+coredumpy.patch_except(directory='./dumps')
+# Create a dump in "./dumps" when there's a unittest failure/error
+coredumpy.patch_unittest(directory='./dumps')
+```
+
+Or you can dump the current frame stack manually
 
 ```python
 import coredumpy
@@ -32,15 +42,6 @@ coredumpy.dump(path=lambda: f"coredumpy_{time.time()}.dump")
 coredumpy.dump(directory='./dumps')
 ```
 
-You can hook the exception so a dump will be automatically created if your program crashes due to an exception
-
-```python
-import coredumpy
-coredumpy.patch_except()
-# patch_except takes the same path/directory arguments as dump
-# coredumpy.patch_except(directory='./dumps')
-```
-
 ### load
 
 Load your dump with

src/coredumpy/__init__.py

@@ -8,6 +8,7 @@
 from .except_hook import patch_except
 from .main import main
 from .type_support import add_supports
+from .unittest_hook import patch_unittest
 
 add_supports()
 
@@ -17,4 +18,5 @@
     "load",
     "main",
     "patch_except",
+    "patch_unittest",
 ]

src/coredumpy/coredumpy.py

@@ -9,6 +9,7 @@
 import pdb
 import tokenize
 import types
+import warnings
 from typing import Callable, Optional, Union
 
 from .patch import patch_all
@@ -48,18 +49,26 @@ def dump(cls,
             if filename not in files:
                 files.add(filename)
 
-            PyObjectProxy.add_object(frame)
+            with warnings.catch_warnings():
+                warnings.simplefilter("ignore")
+                PyObjectProxy.add_object(frame)
             frame = frame.f_back
 
         output_file = get_dump_filename(curr_frame, path, directory)
 
+        file_lines = {}
+
+        for filename in files:
+            if os.path.exists(filename):
+                with tokenize.open(filename) as f:
+                    file_lines[filename] = f.readlines()
+
+        os.makedirs(os.path.dirname(output_file), exist_ok=True)
         with open(output_file, "w") as f:
             json.dump({
                 "objects": PyObjectProxy._objects,
                 "frame": str(id(curr_frame)),
-                "files": {filename: tokenize.open(filename).readlines()
-                          for filename in files
-                          if os.path.exists(filename)}
+                "files": file_lines,
             }, f)
 
         PyObjectProxy.clear()

src/coredumpy/unittest_hook.py

@@ -0,0 +1,50 @@
+# Licensed under the Apache License: http://www.apache.org/licenses/LICENSE-2.0
+# For details: https://github.com/gaogaotiantian/coredumpy/blob/master/NOTICE.txt
+
+
+import unittest
+from typing import Callable, Optional, Union
+
+from .coredumpy import dump
+
+
+def patch_unittest(path: Optional[Union[str, Callable[[], str]]] = None,
+                   directory: Optional[str] = None):
+    """ Patch unittest to coredump when a test fails/raises an exception.
+
+        @param path:
+            The path to save the dump file. It could be a string or a callable that returns a string.
+            if not specified, the default filename will be used
+        @param directory:
+            The directory to save the dump file, only works when path is not specified.
+    """
+
+    _original_addError = unittest.TestResult.addError
+    _original_addFailure = unittest.TestResult.addFailure
+
+    def addError(self, test, err):
+        tb = err[2]
+        while tb.tb_next:
+            tb = tb.tb_next
+        try:
+            filename = dump(tb.tb_frame, path=path, directory=directory)
+            print(f'Your frame stack has been dumped to "{filename}", '
+                  f'open it with\ncoredumpy load {filename}')
+        except Exception:  # pragma: no cover
+            pass
+        _original_addError(self, test, err)
+
+    def addFailure(self, test, err):
+        tb = err[2]
+        while tb.tb_next:
+            tb = tb.tb_next
+        try:
+            filename = dump(tb.tb_frame, path=path, directory=directory)
+            print(f'Your frame stack has been dumped to "{filename}", '
+                  f'open it with\ncoredumpy load {filename}')
+        except Exception:  # pragma: no cover
+            pass
+        _original_addFailure(self, test, err)
+
+    unittest.TestResult.addError = addError  # type: ignore
+    unittest.TestResult.addFailure = addFailure  # type: ignore

src/coredumpy/utils.py

@@ -4,6 +4,7 @@
 
 import datetime
 import os
+import re
 
 
 def get_dump_filename(frame, path, directory):
@@ -14,9 +15,12 @@ def get_dump_filename(frame, path, directory):
             return os.path.abspath(path())
         return os.path.abspath(path)
 
-    funcname = os.path.basename(frame.f_code.co_name)
+    funcname = frame.f_code.co_name
+    lineno = frame.f_lineno
+    funcname = re.sub(r"[^a-zA-Z0-9_]", "", funcname)
+
     d = datetime.datetime.now()
-    filename = f"coredumpy_{funcname}_{d.strftime('%Y%m%d_%H%M%S_%f')}.dump"
+    filename = f"coredumpy_{funcname}_{lineno}_{d.strftime('%Y%m%d_%H%M%S_%f')}.dump"
 
     if directory is None:
         return os.path.abspath(filename)

tests/base.py

@@ -4,6 +4,7 @@
 
 import os
 import subprocess
+import sys
 import tempfile
 import textwrap
 import unittest
@@ -17,7 +18,7 @@ def run_test(self, script, dumppath, commands):
         with tempfile.TemporaryDirectory() as tmpdir:
             with open(f"{tmpdir}/script.py", "w") as f:
                 f.write(script)
-            subprocess.run(normalize_commands(["python", f"{tmpdir}/script.py"]),
+            subprocess.run(normalize_commands([sys.executable, f"{tmpdir}/script.py"]),
                            stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
 
             process = subprocess.Popen(normalize_commands(["coredumpy", "load", dumppath]),
@@ -32,14 +33,16 @@ def run_test(self, script, dumppath, commands):
             pass
         return stdout, stderr
 
-    def run_script(self, script):
+    def run_script(self, script, expected_returncode=0):
         script = textwrap.dedent(script)
         with tempfile.TemporaryDirectory() as tmpdir:
             with open(f"{tmpdir}/script.py", "w") as f:
                 f.write(script)
-            process = subprocess.Popen(normalize_commands(["python", f"{tmpdir}/script.py"]),
+            process = subprocess.Popen(normalize_commands([sys.executable, f"{tmpdir}/script.py"]),
                                        stdout=subprocess.PIPE, stderr=subprocess.PIPE)
             stdout, stderr = process.communicate()
             stdout = stdout.decode()
             stderr = stderr.decode()
+            self.assertEqual(process.returncode, expected_returncode,
+                             f"script failed with return code {process.returncode}\n{stderr}")
         return stdout, stderr

tests/test_basic.py

@@ -2,6 +2,9 @@
 # For details: https://github.com/gaogaotiantian/coredumpy/blob/master/NOTICE.txt
 
 
+import os
+import tempfile
+
 from .base import TestBase
 
 
@@ -33,6 +36,19 @@ def f():
         self.assertIn("[3, {'a': [4, None]}]", stdout)
         self.assertIn("142857", stdout)
 
+    def test_directory(self):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            child_dir = os.path.join(tmpdir, "child")
+            script = f"""
+                import coredumpy
+                coredumpy.dump(directory={repr(tmpdir)})
+                coredumpy.dump(directory={repr(tmpdir)})
+                coredumpy.dump(directory={repr(child_dir)})
+            """
+            self.run_script(script)
+            self.assertEqual(len(os.listdir(tmpdir)), 3)
+            self.assertEqual(len(os.listdir(child_dir)), 1)
+
     def test_except(self):
         script = """
             import coredumpy

tests/test_unittest.py

@@ -0,0 +1,36 @@
+# Licensed under the Apache License: http://www.apache.org/licenses/LICENSE-2.0
+# For details: https://github.com/gaogaotiantian/coredumpy/blob/master/NOTICE.txt
+
+
+import os
+import tempfile
+
+
+from .base import TestBase
+
+
+class TestUnittest(TestBase):
+    def test_unittest_basic(self):
+        with tempfile.TemporaryDirectory() as tempdir:
+            script = f"""
+                import unittest
+                from coredumpy import patch_unittest
+                patch_unittest(directory={repr(tempdir)})
+                class TestUnittest(unittest.TestCase):
+                    def test_bool(self):
+                        self.assertTrue(False)
+                    def test_eq(self):
+                        self.assertEqual(1, 2)
+                    def test_pass(self):
+                        self.assertEqual(1, 1)
+                    def test_error(self):
+                        raise ValueError()
+                unittest.main()
+            """
+            stdout, stderr = self.run_script(script, expected_returncode=1)
+            self.assertIn("FAIL: test_bool", stderr)
+            self.assertIn("FAIL: test_eq", stderr)
+            self.assertIn("ERROR: test_error", stderr)
+            self.assertNotIn("test_pass", stderr)
+            self.assertEqual(stdout.count(tempdir), 6)
+            self.assertEqual(len(os.listdir(tempdir)), 3)

tests/test_utils.py

@@ -8,20 +8,23 @@
 from .base import TestBase
 
 
-class TestUtils(TestBase):
-    def test_get_dump_filename(self):
-        class FakeFrame:
-            def __init__(self, name):
-                self.f_code = FakeCode(name)
+class FakeCode:
+    def __init__(self, name):
+        self.co_name = name
+
+
+class FakeFrame:
+    def __init__(self, name):
+        self.f_code = FakeCode(name)
+        self.f_lineno = 1
 
-        class FakeCode:
-            def __init__(self, name):
-                self.co_name = name
 
+class TestUtils(TestBase):
+    def test_get_dump_filename(self):
         frame = FakeFrame("test_get_dump_filename")
         filename = get_dump_filename(frame, None, None)
         self.assertEqual(filename, os.path.abspath(filename))
-        self.assertIn("test_get_dump_filename", filename)
+        self.assertIn("test_get_dump_filename_1", filename)
 
         filename = get_dump_filename(frame, "test.dump", None)
         self.assertEqual(filename, os.path.abspath("test.dump"))
@@ -31,8 +34,14 @@ def __init__(self, name):
 
         filename = get_dump_filename(frame, None, "dir")
         self.assertEqual(filename, os.path.abspath(filename))
-        self.assertIn("test_get_dump_filename", filename)
+        self.assertIn("test_get_dump_filename_1", filename)
         self.assertIn("dir", filename)
 
         with self.assertRaises(ValueError):
             filename = get_dump_filename(frame, "test.dump", "dir")
+
+    def test_escape_name(self):
+        frame = FakeFrame("<module>")
+        filename = get_dump_filename(frame, None, None)
+        self.assertNotIn("<", filename)
+        self.assertNotIn(">", filename)

tests/util.py

@@ -3,11 +3,12 @@
 
 
 import os
+import sys
 
 
 def normalize_commands(commands):
     if os.getenv("COVERAGE_RUN"):
-        if commands[0] == "python":
+        if commands[0] == "python" or commands[0] == sys.executable:
             commands = ["coverage", "run", "--parallel-mode"] + commands[1:]
         elif commands[0] == "coredumpy":
             commands = ["coverage", "run", "--parallel-mode", "-m", "coredumpy"] + commands[1:]