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 |
|
6 import multiprocessing |
|
7 import multiprocessing.process |
|
8 import sys |
|
9 |
|
10 # An attribute that will be set on modules to indicate that they have been |
|
11 # monkey-patched. |
|
12 PATCHED_MARKER = "_coverage$patched" |
|
13 |
|
14 if sys.version_info >= (3, 4): |
|
15 klass = multiprocessing.process.BaseProcess |
|
16 else: |
|
17 klass = multiprocessing.Process |
|
18 |
|
19 original_bootstrap = klass._bootstrap |
|
20 |
|
21 |
|
22 class ProcessWithCoverage(klass): |
|
23 """A replacement for multiprocess.Process that starts coverage.""" |
|
24 def _bootstrap(self): |
|
25 """Wrapper around _bootstrap to start coverage.""" |
|
26 from coverage import Coverage |
|
27 cov = Coverage(data_suffix=True) |
|
28 cov.start() |
|
29 try: |
|
30 return original_bootstrap(self) |
|
31 finally: |
|
32 cov.stop() |
|
33 cov.save() |
|
34 |
|
35 |
|
36 class Stowaway(object): |
|
37 """An object to pickle, so when it is unpickled, it can apply the monkey-patch.""" |
|
38 def __getstate__(self): |
|
39 return {} |
|
40 |
|
41 def __setstate__(self, state_unused): |
|
42 patch_multiprocessing() |
|
43 |
|
44 |
|
45 def patch_multiprocessing(): |
|
46 """Monkey-patch the multiprocessing module. |
|
47 |
|
48 This enables coverage measurement of processes started by multiprocessing. |
|
49 This is wildly experimental! |
|
50 |
|
51 """ |
|
52 if hasattr(multiprocessing, PATCHED_MARKER): |
|
53 return |
|
54 |
|
55 if sys.version_info >= (3, 4): |
|
56 klass._bootstrap = ProcessWithCoverage._bootstrap |
|
57 else: |
|
58 multiprocessing.Process = ProcessWithCoverage |
|
59 |
|
60 # When spawning processes rather than forking them, we have no state in the |
|
61 # new process. We sneak in there with a Stowaway: we stuff one of our own |
|
62 # objects into the data that gets pickled and sent to the sub-process. When |
|
63 # the Stowaway is unpickled, it's __setstate__ method is called, which |
|
64 # re-applies the monkey-patch. |
|
65 # Windows only spawns, so this is needed to keep Windows working. |
|
66 try: |
|
67 from multiprocessing import spawn # pylint: disable=no-name-in-module |
|
68 original_get_preparation_data = spawn.get_preparation_data |
|
69 except (ImportError, AttributeError): |
|
70 pass |
|
71 else: |
|
72 def get_preparation_data_with_stowaway(name): |
|
73 """Get the original preparation data, and also insert our stowaway.""" |
|
74 d = original_get_preparation_data(name) |
|
75 d['stowaway'] = Stowaway() |
|
76 return d |
|
77 |
|
78 spawn.get_preparation_data = get_preparation_data_with_stowaway |
|
79 |
|
80 setattr(multiprocessing, PATCHED_MARKER, True) |
|
81 |
|
82 # |
|
83 # eflag: FileType = Python2 |
|