9 |
9 |
10 # An attribute that will be set on modules to indicate that they have been |
10 # An attribute that will be set on modules to indicate that they have been |
11 # monkey-patched. |
11 # monkey-patched. |
12 PATCHED_MARKER = "_coverage$patched" |
12 PATCHED_MARKER = "_coverage$patched" |
13 |
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 |
14 |
44 |
15 def patch_multiprocessing(): |
45 def patch_multiprocessing(): |
16 """Monkey-patch the multiprocessing module. |
46 """Monkey-patch the multiprocessing module. |
17 |
47 |
18 This enables coverage measurement of processes started by multiprocessing. |
48 This enables coverage measurement of processes started by multiprocessing. |
21 """ |
51 """ |
22 if hasattr(multiprocessing, PATCHED_MARKER): |
52 if hasattr(multiprocessing, PATCHED_MARKER): |
23 return |
53 return |
24 |
54 |
25 if sys.version_info >= (3, 4): |
55 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 |
56 klass._bootstrap = ProcessWithCoverage._bootstrap |
47 else: |
57 else: |
48 multiprocessing.Process = ProcessWithCoverage |
58 multiprocessing.Process = ProcessWithCoverage |
49 |
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 |
50 setattr(multiprocessing, PATCHED_MARKER, True) |
80 setattr(multiprocessing, PATCHED_MARKER, True) |