From 17d551e25895b1ccea5f4c7e1602eef58c2829a0 Mon Sep 17 00:00:00 2001
From: Tian Gao <gaogaotiantian@hotmail.com>
Date: Tue, 30 Apr 2024 21:51:09 -0700
Subject: [PATCH] Add run command (#25)

---
 .coveragerc                      |   2 +-
 README.md                        |  26 +++---
 src/coredumpy/coredumpy.py       | 131 ++++++++++++++++++++++++++++---
 src/coredumpy/main.py            |  14 +++-
 src/coredumpy/py_object_proxy.py |   3 +
 tests/__init__.py                |   2 +
 tests/base.py                    |  18 ++++-
 tests/data/__init__.py           |   2 +
 tests/data/failed.py             |  19 +++++
 tests/test_basic.py              |  34 ++++++++
 tests/test_unittest.py           |  12 ++-
 tests/util.py                    |   4 +-
 12 files changed, 239 insertions(+), 28 deletions(-)
 create mode 100644 tests/data/__init__.py
 create mode 100644 tests/data/failed.py

diff --git a/.coveragerc b/.coveragerc
index fe37262..e0d692b 100644
--- a/.coveragerc
+++ b/.coveragerc
@@ -1,3 +1,3 @@
 [run]
 cover_pylib = True
-source = coredumpy
+source_pkgs = coredumpy
diff --git a/README.md b/README.md
index 9dedce9..551a8e1 100644
--- a/README.md
+++ b/README.md
@@ -15,9 +15,24 @@ coredumpy saves your crash site for post-mortem debugging.
 
 ### dump
 
-In most cases, you only need to hook `coredumpy` to some triggers
+For `pytest`, you can use `coredumpy` as a plugin
 
-For `Exception` and `unittest`, patch with a simple line
+```
+# Create a dump in "./dumps" when there's a pytest failure/error
+pytest --enable-coredumpy --coredumpy-dir ./dumps
+```
+
+For `Exception` and `unittest`, you can use `coredumpy run` command.
+A dump will be generated when there's an unhandled exception or a test failure
+
+```
+# with no argument coredumpy run will generate the dump in the current dir
+coredumpy run my_script.py
+coredumpy run my_script.py --directory ./dumps
+coredumpy run -m unittest --directory ./dumps
+```
+
+Or you can patch explicitly in your code and execute the script/module as usual
 
 ```python
 import coredumpy
@@ -27,13 +42,6 @@ coredumpy.patch_except(directory='./dumps')
 coredumpy.patch_unittest(directory='./dumps')
 ```
 
-For `pytest`, you can use `coredumpy` as a plugin
-
-```
-# Create a dump in "./dumps" when there's a pytest failure/error
-pytest --enable-coredumpy --coredumpy-dir ./dumps
-```
-
 <details>
 
 <summary>
diff --git a/src/coredumpy/coredumpy.py b/src/coredumpy/coredumpy.py
index f0e75de..2e3ec78 100644
--- a/src/coredumpy/coredumpy.py
+++ b/src/coredumpy/coredumpy.py
@@ -10,10 +10,12 @@
 import os
 import pdb
 import platform
+import sys
 import tokenize
 import textwrap
 import types
 import warnings
+from types import CodeType
 from typing import Callable, Optional, Union
 
 from .patch import patch_all
@@ -21,6 +23,85 @@
 from .utils import get_dump_filename
 
 
+class _ExecutableTarget:
+    filename: str
+    code: Union[CodeType, str]
+    namespace: dict
+
+
+class _ScriptTarget(_ExecutableTarget):
+    def __init__(self, target):
+        self._target = os.path.realpath(target)
+
+        if not os.path.exists(self._target):
+            print(f'Error: {target} does not exist')
+            sys.exit(1)
+        if os.path.isdir(self._target):
+            print(f'Error: {target} is a directory')
+            sys.exit(1)
+
+        # If safe_path(-P) is not set, sys.path[0] is the directory
+        # of coredumpy, and we should replace it with the directory of the script
+        if not getattr(sys.flags, "safe_path", None):
+            sys.path[0] = os.path.dirname(self._target)
+
+    @property
+    def filename(self):
+        return self._target
+
+    @property
+    def code(self):
+        # Open the file each time because the file may be modified
+        import io
+        with io.open_code(self._target) as fp:
+            return f"exec(compile({fp.read()!r}, {self._target!r}, 'exec'))"
+
+    @property
+    def namespace(self):
+        return dict(
+            __name__='__main__',
+            __file__=self._target,
+            __builtins__=__builtins__,
+            __spec__=None,
+        )
+
+
+class _ModuleTarget(_ExecutableTarget):
+    def __init__(self, target):
+        self._target = target
+
+        import runpy
+        try:
+            sys.path.insert(0, os.getcwd())
+            _, self._spec, self._code = runpy._get_module_details(self._target)
+        except ImportError as e:
+            print(f"ImportError: {e}")
+            sys.exit(1)
+        except Exception:  # pragma: no cover
+            import traceback
+            traceback.print_exc()
+            sys.exit(1)
+
+    @property
+    def filename(self):
+        return self._code.co_filename
+
+    @property
+    def code(self):
+        return self._code
+
+    @property
+    def namespace(self):
+        return dict(
+            __name__='__main__',
+            __file__=os.path.normcase(os.path.abspath(self.filename)),
+            __package__=self._spec.parent,
+            __loader__=self._spec.loader,
+            __spec__=self._spec,
+            __builtins__=__builtins__,
+        )
+
+
 class Coredumpy:
     @classmethod
     def dump(cls,
@@ -47,20 +128,20 @@ def dump(cls,
             inner_frame = inspect.currentframe()
             assert inner_frame is not None
             frame = inner_frame.f_back
-        curr_frame = frame
+
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+            PyObjectProxy.add_object(frame)
+
+        output_file = get_dump_filename(frame, path, directory)
+        frame_id = str(id(frame))
+
         while frame:
             filename = frame.f_code.co_filename
-
             if filename not in files:
                 files.add(filename)
-
-            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:
@@ -73,7 +154,7 @@ def dump(cls,
         with gzip.open(output_file, "wt") as f:
             json.dump({
                 "objects": PyObjectProxy._objects,
-                "frame": str(id(curr_frame)),
+                "frame": frame_id,
                 "files": file_lines,
                 "description": description,
                 "metadata": cls.get_metadata()
@@ -121,6 +202,37 @@ def peek(cls, path):
         if data["description"]:
             print(textwrap.indent(data["description"], "    "))
 
+    @classmethod
+    def run(cls, options):
+        if options.module:
+            file = options.module
+            target = _ModuleTarget(file)
+        else:
+            if not options.args:
+                print("Error: no script specified")
+                sys.exit(1)
+            file = options.args.pop(0)
+            target = _ScriptTarget(file)
+
+        sys.argv[:] = [file] + options.args
+
+        import __main__
+        __main__.__dict__.clear()
+        __main__.__dict__.update(target.namespace)
+
+        cmd = target.code
+
+        if isinstance(cmd, str):
+            cmd = compile(cmd, "<string>", "exec")
+
+        from .except_hook import patch_except
+        patch_except(path=options.path, directory=options.directory)
+
+        from .unittest_hook import patch_unittest
+        patch_unittest(path=options.path, directory=options.directory)
+
+        exec(cmd, __main__.__dict__, __main__.__dict__)
+
     @classmethod
     def get_metadata(cls):
         from coredumpy import __version__
@@ -140,3 +252,4 @@ def get_metadata(cls):
 dump = Coredumpy.dump
 load = Coredumpy.load
 peek = Coredumpy.peek
+run = Coredumpy.run
diff --git a/src/coredumpy/main.py b/src/coredumpy/main.py
index 69e63aa..81d33bb 100644
--- a/src/coredumpy/main.py
+++ b/src/coredumpy/main.py
@@ -5,18 +5,24 @@
 import argparse
 import os
 
-from .coredumpy import load, peek
+from .coredumpy import load, peek, run
 
 
 def main():
     parser = argparse.ArgumentParser()
     subparsers = parser.add_subparsers(dest="command")
 
+    subparsers_run = subparsers.add_parser("run", help="Run a file/module with coredumpy enabled.")
+    subparsers_run.add_argument("-m", metavar="module", dest="module")
+    subparsers_run.add_argument("--path", help="The path of dump file", default=None)
+    subparsers_run.add_argument("--directory", help="The directory of dump file", default=None)
+    subparsers_run.add_argument("args", nargs="*")
+
     subparsers_load = subparsers.add_parser("load", help="Load a dump file.")
     subparsers_load.add_argument("file", type=str, help="The dump file to load.")
 
-    subparsers_load = subparsers.add_parser("peek", help="Peek a dump file.")
-    subparsers_load.add_argument("files", help="The dump file to load.", nargs="+")
+    subparsers_peek = subparsers.add_parser("peek", help="Peek a dump file.")
+    subparsers_peek.add_argument("files", help="The dump file to load.", nargs="+")
 
     args = parser.parse_args()
 
@@ -42,3 +48,5 @@ def main():
                         pass
             else:
                 print(f"File {file} not found.")
+    elif args.command == "run":
+        run(args)
diff --git a/src/coredumpy/py_object_proxy.py b/src/coredumpy/py_object_proxy.py
index e88953a..bc0d360 100644
--- a/src/coredumpy/py_object_proxy.py
+++ b/src/coredumpy/py_object_proxy.py
@@ -53,6 +53,9 @@ def _add_object(cls, obj):
         id_str = str(id(obj))
         if id_str not in cls._objects:
             cls._objects[id_str] = {"type": "_coredumpy_unknown"}
+            if obj is cls._objects or obj is cls._pending_objects:
+                # Avoid changing the dict while dumping
+                return
             cls._pending_objects.put((cls._current_recursion_depth + 1, obj))
 
     @classmethod
diff --git a/tests/__init__.py b/tests/__init__.py
index e69de29..a92343d 100644
--- a/tests/__init__.py
+++ b/tests/__init__.py
@@ -0,0 +1,2 @@
+# Licensed under the Apache License: http://www.apache.org/licenses/LICENSE-2.0
+# For details: https://github.com/gaogaotiantian/coredumpy/blob/master/NOTICE.txt
diff --git a/tests/base.py b/tests/base.py
index ade24ed..81eaa60 100644
--- a/tests/base.py
+++ b/tests/base.py
@@ -13,13 +13,17 @@
 
 
 class TestBase(unittest.TestCase):
-    def run_test(self, script, dumppath, commands):
+    def run_test(self, script, dumppath, commands, use_cli_run=False):
         script = textwrap.dedent(script)
         with tempfile.TemporaryDirectory() as tmpdir:
             with open(f"{tmpdir}/script.py", "w") as f:
                 f.write(script)
-            subprocess.run(normalize_commands([sys.executable, f"{tmpdir}/script.py"]),
-                           stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
+            if use_cli_run:
+                subprocess.run(normalize_commands(["coredumpy", "run", f"{tmpdir}/script.py", "--path", dumppath]),
+                               stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
+            else:
+                subprocess.run(normalize_commands([sys.executable, f"{tmpdir}/script.py"]),
+                               stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
 
             process = subprocess.Popen(normalize_commands(["coredumpy", "load", dumppath]),
                                        stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
@@ -47,6 +51,14 @@ def run_script(self, script, expected_returncode=0):
                              f"script failed with return code {process.returncode}\n{stderr}")
         return stdout, stderr
 
+    def run_run(self, args):
+        process = subprocess.Popen(normalize_commands(["coredumpy", "run"] + args),
+                                   stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+        stdout, stderr = process.communicate()
+        stdout = stdout.decode(errors='backslashreplace')
+        stderr = stderr.decode(errors='backslashreplace')
+        return stdout, stderr
+
     def run_peek(self, paths):
         process = subprocess.Popen(normalize_commands(["coredumpy", "peek"] + paths),
                                    stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
diff --git a/tests/data/__init__.py b/tests/data/__init__.py
new file mode 100644
index 0000000..a92343d
--- /dev/null
+++ b/tests/data/__init__.py
@@ -0,0 +1,2 @@
+# Licensed under the Apache License: http://www.apache.org/licenses/LICENSE-2.0
+# For details: https://github.com/gaogaotiantian/coredumpy/blob/master/NOTICE.txt
diff --git a/tests/data/failed.py b/tests/data/failed.py
new file mode 100644
index 0000000..5ea538f
--- /dev/null
+++ b/tests/data/failed.py
@@ -0,0 +1,19 @@
+# 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
+
+
+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()
diff --git a/tests/test_basic.py b/tests/test_basic.py
index 5d545c8..efabc84 100644
--- a/tests/test_basic.py
+++ b/tests/test_basic.py
@@ -67,6 +67,40 @@ def g(arg):
         self.assertIn("return 1 / arg", stdout)
         self.assertIn("0", stdout)
 
+    def test_cli(self):
+        script = """
+            def g(arg):
+                return 1 / arg
+            def f(x):
+                a = 142857
+                g(x)
+            f(0)
+        """
+        stdout, _ = self.run_test(script, "coredumpy_dump", [
+            "w",
+            "p arg",
+            "u",
+            "p a",
+            "q"
+        ], use_cli_run=True)
+
+        self.assertIn("-> f(0)", stdout)
+        self.assertIn("-> g(x)", stdout)
+        self.assertIn("142857", stdout)
+
+    def test_cli_invalid(self):
+        stdout, _ = self.run_run([])
+        self.assertIn("Error", stdout)
+
+        stdout, _ = self.run_run(["notexist.py"])
+        self.assertIn("Error", stdout)
+
+        stdout, _ = self.run_run([os.path.dirname(__file__)])
+        self.assertIn("Error", stdout)
+
+        stdout, _ = self.run_run(["-m", "nonexistmodule"])
+        self.assertIn("Error", stdout)
+
     def test_peek(self):
         with tempfile.TemporaryDirectory() as tmpdir:
             script = f"""
diff --git a/tests/test_unittest.py b/tests/test_unittest.py
index 612f64c..ef81099 100644
--- a/tests/test_unittest.py
+++ b/tests/test_unittest.py
@@ -5,7 +5,6 @@
 import os
 import tempfile
 
-
 from .base import TestBase
 
 
@@ -34,3 +33,14 @@ def test_error(self):
             self.assertNotIn("test_pass", stderr)
             self.assertEqual(stdout.count(tempdir), 3)
             self.assertEqual(len(os.listdir(tempdir)), 3)
+
+    def test_unittest_with_cli(self):
+        with tempfile.TemporaryDirectory() as tempdir:
+            stdout, stderr = self.run_run(["-m", "unittest", "tests.data.failed",
+                                           "--directory", tempdir])
+            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), 3)
+            self.assertEqual(len(os.listdir(tempdir)), 3)
diff --git a/tests/util.py b/tests/util.py
index 7b5a78c..101d47f 100644
--- a/tests/util.py
+++ b/tests/util.py
@@ -9,7 +9,7 @@
 def normalize_commands(commands):
     if os.getenv("COVERAGE_RUN"):
         if commands[0] == "python" or commands[0] == sys.executable:
-            commands = ["coverage", "run", "--parallel-mode"] + commands[1:]
+            commands = [sys.executable, "-m", "coverage", "run", "--parallel-mode"] + commands[1:]
         elif commands[0] == "coredumpy":
-            commands = ["coverage", "run", "--parallel-mode", "-m", "coredumpy"] + commands[1:]
+            commands = [sys.executable, "-m", "coverage", "run", "--parallel-mode", "-m", "coredumpy"] + commands[1:]
     return commands