feat: patch os._exit #310

Author: nedbat

coverage/config.py

@@ -207,6 +207,7 @@ def __init__(self) -> None:
         self.disable_warnings: list[str] = []
         self.dynamic_context: str | None = None
         self.parallel = False
+        self.patch: list[str] = []
         self.plugins: list[str] = []
         self.relative_files = False
         self.run_include: list[str] = []
@@ -267,6 +268,7 @@ def __init__(self) -> None:
         "debug", "concurrency", "plugins",
         "report_omit", "report_include",
         "run_omit", "run_include",
+        "patch",
     }
 
     def from_args(self, **kwargs: TConfigValueIn) -> None:
@@ -390,6 +392,7 @@ def copy(self) -> CoverageConfig:
         ("disable_warnings", "run:disable_warnings", "list"),
         ("dynamic_context", "run:dynamic_context"),
         ("parallel", "run:parallel", "boolean"),
+        ("patch", "run:patch", "list"),
         ("plugins", "run:plugins", "list"),
         ("relative_files", "run:relative_files", "boolean"),
         ("run_include", "run:include", "list"),

coverage/control.py

@@ -42,6 +42,7 @@
 from coverage.misc import bool_or_none, join_regex
 from coverage.misc import DefaultValue, ensure_dir_for_file, isolate_module
 from coverage.multiproc import patch_multiprocessing
+from coverage.patch import apply_patches
 from coverage.plugin import FileReporter
 from coverage.plugin_support import Plugins, TCoverageInit
 from coverage.python import PythonFileReporter
@@ -677,6 +678,8 @@ def start(self) -> None:
         if self._auto_load:
             self.load()
 
+        apply_patches(self, self.config)
+
         self._collector.start()
         self._started = True
         self._instances.append(self)

coverage/patch.py

@@ -0,0 +1,27 @@
+# Licensed under the Apache License: http://www.apache.org/licenses/LICENSE-2.0
+# For details: https://github.com/nedbat/coveragepy/blob/master/NOTICE.txt
+
+"""Invasive patches for coverage.py"""
+
+from __future__ import annotations
+
+import os
+
+from typing import NoReturn, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from coverage import Coverage
+    from coverage.config import CoverageConfig
+
+def apply_patches(cov: Coverage, config: CoverageConfig) -> None:
+    """Apply invasive patches requested by `[run] patch=`."""
+
+    if "os._exit" in config.patch:
+        _old_os_exit = os._exit
+        def _coverage_os_exit_patch(status: int) -> NoReturn:
+            try:
+                cov.save()
+            except: # pylint: disable=bare-except
+                pass
+            _old_os_exit(status)
+        os._exit = _coverage_os_exit_patch

tests/test_process.py

@@ -421,6 +421,43 @@ def test_fork(self) -> None:
         reported_pids = {line.split(".")[0] for line in debug_text.splitlines()}
         assert len(reported_pids) == 2
 
+    @pytest.mark.skipif(not hasattr(os, "fork"), reason="Can't test os.fork, it doesn't exist.")
+    @pytest.mark.parametrize("patch", [False, True])
+    def test_os_exit(self, patch: bool) -> None:
+        self.make_file("forky.py", """\
+            import os
+            import tempfile
+            import time
+
+            complete_file = tempfile.mkstemp()[1]
+            pid = os.fork()
+            if pid:
+                while pid: # 3.9 wouldn't count "while True": change this. PYVERSION
+                    with open(complete_file, encoding="ascii") as f:
+                        data = f.read()
+                    if "Complete" in data:
+                        break
+                    time.sleep(.02)
+                os.remove(complete_file)
+            else:
+                time.sleep(.1)
+                with open(complete_file, mode="w", encoding="ascii") as f:
+                    f.write("Complete")
+                os._exit(0)
+            """)
+        total_lines = 17
+        if patch:
+            self.make_file(".coveragerc", "[run]\npatch = os._exit\n")
+        self.run_command("coverage run -p forky.py")
+        self.run_command("coverage combine")
+        data = coverage.CoverageData()
+        data.read()
+        seen = line_counts(data)["forky.py"]
+        if patch:
+            assert seen == total_lines
+        else:
+            assert seen < total_lines
+
     def test_warnings_during_reporting(self) -> None:
         # While fixing issue #224, the warnings were being printed far too
         # often.  Make sure they're not any more.