1# Licensed under the Apache License: http://www.apache.org/licenses/LICENSE-2.0
2# For details: https://bitbucket.org/ned/coveragepy/src/default/NOTICE.txt
3
4"""Monkey-patching to make coverage.py work right in some cases."""
5
6import multiprocessing
7import multiprocessing.process
8import sys
9
10# An attribute that will be set on modules to indicate that they have been
11# monkey-patched.
12PATCHED_MARKER = "_coverage$patched"
13
14
15def patch_multiprocessing():
16    """Monkey-patch the multiprocessing module.
17
18    This enables coverage measurement of processes started by multiprocessing.
19    This is wildly experimental!
20
21    """
22    if hasattr(multiprocessing, PATCHED_MARKER):
23        return
24
25    if sys.version_info >= (3, 4):
26        klass = multiprocessing.process.BaseProcess
27    else:
28        klass = multiprocessing.Process
29
30    original_bootstrap = klass._bootstrap
31
32    class ProcessWithCoverage(klass):
33        """A replacement for multiprocess.Process that starts coverage."""
34        def _bootstrap(self):
35            """Wrapper around _bootstrap to start coverage."""
36            from coverage import Coverage
37            cov = Coverage(data_suffix=True)
38            cov.start()
39            try:
40                return original_bootstrap(self)
41            finally:
42                cov.stop()
43                cov.save()
44
45    if sys.version_info >= (3, 4):
46        klass._bootstrap = ProcessWithCoverage._bootstrap
47    else:
48        multiprocessing.Process = ProcessWithCoverage
49
50    setattr(multiprocessing, PATCHED_MARKER, True)
51