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 6import multiprocessing 7import multiprocessing.process 8import os 9import os.path 10import sys 11import traceback 12 13from coverage import env 14from coverage.misc import contract 15 16# An attribute that will be set on the module to indicate that it has been 17# monkey-patched. 18PATCHED_MARKER = "_coverage$patched" 19 20 21if env.PYVERSION >= (3, 4): 22 OriginalProcess = multiprocessing.process.BaseProcess 23else: 24 OriginalProcess = multiprocessing.Process 25 26original_bootstrap = OriginalProcess._bootstrap 27 28class 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 56class 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) 69def 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) 112