From 0fd58f32e09bdef8eeea7f62a63546883ac6147d Mon Sep 17 00:00:00 2001
From: Brian Wickman <wickman@twitter.com>
Date: Thu, 16 Apr 2015 11:47:14 -0700
Subject: [PATCH] Unify coverage and profiling handling and semantics.

---
 pex/pex.py | 85 +++++++++++++++++++++++++++++++++++-------------------
 1 file changed, 56 insertions(+), 29 deletions(-)

diff --git a/pex/pex.py b/pex/pex.py
index 94aa14878..412834e79 100644
--- a/pex/pex.py
+++ b/pex/pex.py
@@ -13,7 +13,7 @@
 import pkg_resources
 from pkg_resources import EntryPoint, find_distributions
 
-from .common import die, safe_mkdir
+from .common import die
 from .compatibility import exec_function
 from .environment import PEXEnvironment
 from .finders import get_script_from_distributions
@@ -39,15 +39,6 @@ class PEX(object):  # noqa: T000
   class Error(Exception): pass
   class NotFound(Error): pass
 
-  @classmethod
-  def start_coverage(cls):
-    try:
-      import coverage
-      cov = coverage.coverage(auto_data=True, data_suffix=True)
-      cov.start()
-    except ImportError:
-      sys.stderr.write('Could not bootstrap coverage module!\n')
-
   @classmethod
   def clean_environment(cls, forking=False):
     try:
@@ -234,6 +225,59 @@ def patch_all(path, path_importer_cache, modules):
     finally:
       patch_all(old_sys_path, old_sys_path_importer_cache, old_sys_modules)
 
+  @classmethod
+  def _wrap_coverage(cls, runner, *args):
+    if 'PEX_COVERAGE' not in os.environ and 'PEX_COVERAGE_FILENAME' not in os.environ:
+      runner(*args)
+      return
+
+    try:
+      import coverage
+    except ImportError:
+      die('Could not bootstrap coverage module, aborting.')
+
+    if 'PEX_COVERAGE_FILENAME' in os.environ:
+      cov = coverage.coverage(data_file=os.environ['PEX_COVERAGE_FILENAME'])
+    else:
+      cov = coverage.coverage(data_suffix=True)
+
+    TRACER.log('Starting coverage.')
+    cov.start()
+
+    try:
+      runner(*args)
+    finally:
+      TRACER.log('Stopping coverage')
+      cov.stop()
+
+      # TODO(wickman) Post-process coverage to elide $PEX_ROOT and make
+      # the report more useful/less noisy.
+      if 'PEX_COVERAGE_FILENAME' in os.environ:
+        cov.save()
+      else:
+        cov.report(show_missing=False, ignore_errors=True, file=sys.stdout)
+
+  @classmethod
+  def _wrap_profiling(cls, runner, *args):
+    if 'PEX_PROFILE' not in os.environ and 'PEX_PROFILE_FILENAME' not in os.environ:
+      runner(*args)
+      return
+
+    try:
+      import cProfile as profile
+    except ImportError:
+      import profile
+
+    profiler = profile.Profile()
+
+    try:
+      return profiler.runcall(runner, *args)
+    finally:
+      if 'PEX_PROFILE_FILENAME' in os.environ:
+        profiler.dump_stats(os.environ['PEX_PROFILE_FILENAME'])
+      else:
+        profiler.print_stats(sort=os.environ.get('PEX_PROFILE_SORT', 'cumulative'))
+
   def execute(self):
     """Execute the PEX.
 
@@ -243,14 +287,12 @@ def execute(self):
     try:
       with self.patch_sys():
         working_set = self._env.activate()
-        if 'PEX_COVERAGE' in os.environ:
-          self.start_coverage()
         TRACER.log('PYTHONPATH contains:')
         for element in sys.path:
           TRACER.log('  %c %s' % (' ' if os.path.exists(element) else '*', element))
         TRACER.log('  * - paths that do not exist or will be imported via zipimport')
         with self.patch_pkg_resources(working_set):
-          self._execute()
+          self._wrap_coverage(self._wrap_profiling, self._execute)
     except Exception:
       # Allow the current sys.excepthook to handle this app exception before we tear things down in
       # finally, then reraise so that the exit status is reflected correctly.
@@ -340,25 +382,10 @@ def execute_content(cls, name, content, argv0=None):
         globals().pop('__file__')
       sys.argv[0] = old_argv0
 
-  # TODO(wickman) Find a way to make PEX_PROFILE work with all execute_*
   @classmethod
   def execute_entry(cls, entry_point):
     runner = cls.execute_pkg_resources if ':' in entry_point else cls.execute_module
-
-    if 'PEX_PROFILE' not in os.environ:
-      runner(entry_point)
-    else:
-      import pstats, cProfile
-      profile_output = os.environ['PEX_PROFILE']
-      safe_mkdir(os.path.dirname(profile_output))
-      cProfile.runctx('runner(entry_point)', globals=globals(), locals=locals(),
-                      filename=profile_output)
-      try:
-        entries = int(os.environ.get('PEX_PROFILE_ENTRIES', 1000))
-      except ValueError:
-        entries = 1000
-      pstats.Stats(profile_output).sort_stats(
-          os.environ.get('PEX_PROFILE_SORT', 'cumulative')).print_stats(entries)
+    runner(entry_point)
 
   @staticmethod
   def execute_module(module_name):