|
1 # Licensed under the Apache License: http://www.apache.org/licenses/LICENSE-2.0 |
|
2 # For details: https://github.com/nedbat/coveragepy/blob/master/NOTICE.txt |
|
3 |
|
4 """Monkey-patching to add multiprocessing support for coverage.py""" |
|
5 |
|
6 import multiprocessing |
|
7 import multiprocessing.process |
|
8 import os |
|
9 import os.path |
|
10 import sys |
|
11 import traceback |
|
12 |
|
13 from coverage import env |
|
14 from coverage.misc import contract |
|
15 |
|
16 # An attribute that will be set on the module to indicate that it has been |
|
17 # monkey-patched. |
|
18 PATCHED_MARKER = "_coverage$patched" |
|
19 |
|
20 |
|
21 if env.PYVERSION >= (3, 4): |
|
22 OriginalProcess = multiprocessing.process.BaseProcess |
|
23 else: |
|
24 OriginalProcess = multiprocessing.Process |
|
25 |
|
26 original_bootstrap = OriginalProcess._bootstrap |
|
27 |
|
28 class ProcessWithCoverage(OriginalProcess): # pylint: disable=abstract-method |
|
29 """A replacement for multiprocess.Process that starts coverage.""" |
|
30 |
|
31 def _bootstrap(self, *args, **kwargs): # pylint: disable=signature-differs |
|
32 """Wrapper around _bootstrap to start coverage.""" |
|
33 try: |
|
34 from coverage import Coverage # avoid circular import |
|
35 cov = Coverage(data_suffix=True) |
|
36 cov._warn_preimported_source = False |
|
37 cov.start() |
|
38 debug = cov._debug |
|
39 if debug.should("multiproc"): |
|
40 debug.write("Calling multiprocessing bootstrap") |
|
41 except Exception: |
|
42 print("Exception during multiprocessing bootstrap init:") |
|
43 traceback.print_exc(file=sys.stdout) |
|
44 sys.stdout.flush() |
|
45 raise |
|
46 try: |
|
47 return original_bootstrap(self, *args, **kwargs) |
|
48 finally: |
|
49 if debug.should("multiproc"): |
|
50 debug.write("Finished multiprocessing bootstrap") |
|
51 cov.stop() |
|
52 cov.save() |
|
53 if debug.should("multiproc"): |
|
54 debug.write("Saved multiprocessing data") |
|
55 |
|
56 class Stowaway(object): |
|
57 """An object to pickle, so when it is unpickled, it can apply the monkey-patch.""" |
|
58 def __init__(self, rcfile): |
|
59 self.rcfile = rcfile |
|
60 |
|
61 def __getstate__(self): |
|
62 return {'rcfile': self.rcfile} |
|
63 |
|
64 def __setstate__(self, state): |
|
65 patch_multiprocessing(state['rcfile']) |
|
66 |
|
67 |
|
68 @contract(rcfile=str) |
|
69 def patch_multiprocessing(rcfile): |
|
70 """Monkey-patch the multiprocessing module. |
|
71 |
|
72 This enables coverage measurement of processes started by multiprocessing. |
|
73 This involves aggressive monkey-patching. |
|
74 |
|
75 `rcfile` is the path to the rcfile being used. |
|
76 |
|
77 """ |
|
78 |
|
79 if hasattr(multiprocessing, PATCHED_MARKER): |
|
80 return |
|
81 |
|
82 if env.PYVERSION >= (3, 4): |
|
83 OriginalProcess._bootstrap = ProcessWithCoverage._bootstrap |
|
84 else: |
|
85 multiprocessing.Process = ProcessWithCoverage |
|
86 |
|
87 # Set the value in ProcessWithCoverage that will be pickled into the child |
|
88 # process. |
|
89 os.environ["COVERAGE_RCFILE"] = os.path.abspath(rcfile) |
|
90 |
|
91 # When spawning processes rather than forking them, we have no state in the |
|
92 # new process. We sneak in there with a Stowaway: we stuff one of our own |
|
93 # objects into the data that gets pickled and sent to the sub-process. When |
|
94 # the Stowaway is unpickled, it's __setstate__ method is called, which |
|
95 # re-applies the monkey-patch. |
|
96 # Windows only spawns, so this is needed to keep Windows working. |
|
97 try: |
|
98 from multiprocessing import spawn |
|
99 original_get_preparation_data = spawn.get_preparation_data |
|
100 except (ImportError, AttributeError): |
|
101 pass |
|
102 else: |
|
103 def get_preparation_data_with_stowaway(name): |
|
104 """Get the original preparation data, and also insert our stowaway.""" |
|
105 d = original_get_preparation_data(name) |
|
106 d['stowaway'] = Stowaway(rcfile) |
|
107 return d |
|
108 |
|
109 spawn.get_preparation_data = get_preparation_data_with_stowaway |
|
110 |
|
111 setattr(multiprocessing, PATCHED_MARKER, True) |