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