embed.py 2.5 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697
  1. """Activate coverage at python startup if appropriate.
  2. The python site initialisation will ensure that anything we import
  3. will be removed and not visible at the end of python startup. However
  4. we minimise all work by putting these init actions in this separate
  5. module and only importing what is needed when needed.
  6. For normal python startup when coverage should not be activated the pth
  7. file checks a single env var and does not import or call the init fn
  8. here.
  9. For python startup when an ancestor process has set the env indicating
  10. that code coverage is being collected we activate coverage based on
  11. info passed via env vars.
  12. """
  13. import os
  14. import signal
  15. active_cov = None
  16. def multiprocessing_start(_):
  17. cov = init()
  18. if cov:
  19. multiprocessing.util.Finalize(None, cleanup, args=(cov,), exitpriority=1000)
  20. try:
  21. import multiprocessing.util
  22. except ImportError:
  23. pass
  24. else:
  25. multiprocessing.util.register_after_fork(multiprocessing_start, multiprocessing_start)
  26. def init():
  27. # Only continue if ancestor process has set everything needed in
  28. # the env.
  29. global active_cov
  30. cov_source = os.environ.get('COV_CORE_SOURCE')
  31. cov_config = os.environ.get('COV_CORE_CONFIG')
  32. cov_datafile = os.environ.get('COV_CORE_DATAFILE')
  33. cov_branch = True if os.environ.get('COV_CORE_BRANCH') == 'enabled' else None
  34. if cov_datafile:
  35. # Import what we need to activate coverage.
  36. import coverage
  37. # Determine all source roots.
  38. if cov_source == os.pathsep:
  39. cov_source = None
  40. else:
  41. cov_source = cov_source.split(os.pathsep)
  42. if cov_config == os.pathsep:
  43. cov_config = True
  44. # Activate coverage for this process.
  45. cov = active_cov = coverage.coverage(
  46. source=cov_source,
  47. branch=cov_branch,
  48. data_suffix=True,
  49. config_file=cov_config,
  50. auto_data=True,
  51. data_file=cov_datafile
  52. )
  53. cov.load()
  54. cov.start()
  55. cov._warn_no_data = False
  56. cov._warn_unimported_source = False
  57. return cov
  58. def _cleanup(cov):
  59. if cov is not None:
  60. cov.stop()
  61. cov.save()
  62. def cleanup(cov=None):
  63. global active_cov
  64. _cleanup(cov)
  65. if active_cov is not cov:
  66. _cleanup(active_cov)
  67. active_cov = None
  68. multiprocessing_finish = cleanup # in case someone dared to use this internal
  69. def _sigterm_handler(*_):
  70. cleanup()
  71. def cleanup_on_sigterm():
  72. signal.signal(signal.SIGTERM, _sigterm_handler)