WIP

Author: nedbat

src/pytest_cov/plugin.py

@@ -4,6 +4,7 @@
 import warnings
 
 import pytest
+import coverage
 from coverage.misc import CoverageException
 
 from . import compat
@@ -48,6 +49,16 @@ def validate_fail_under(num_str):
         return float(num_str)
 
 
+def validate_contexts(arg):
+    if coverage.version_info <= (5, 0):
+        msg = '--cov-contexts is only supported with coverage.py >= 5.x'
+        raise argparse.ArgumentTypeError(msg)
+    if arg != "test":
+        msg = '--cov-contexts=test is the only supported value'
+        raise argparse.ArgumentTypeError(msg)
+    return arg
+
+
 class StoreReport(argparse.Action):
     def __call__(self, parser, namespace, values, option_string=None):
         report_type, file = values
@@ -88,6 +99,9 @@ def pytest_addoption(parser):
                          'Default: False')
     group.addoption('--cov-branch', action='store_true', default=None,
                     help='Enable branch coverage.')
+    group.addoption('--cov-contexts', action='store', metavar='CONTEXT',
+                    type=validate_contexts,
+                    help='Dynamic contexts to use. "test" for now.')
 
 
 def _prepare_cov_source(cov_source):
@@ -131,6 +145,7 @@ def __init__(self, options, pluginmanager, start=True):
         self.failed = False
         self._started = False
         self._disabled = False
+        self.test_contexts = False
         self.options = options
 
         is_dist = (getattr(options, 'numprocesses', False) or
@@ -146,6 +161,9 @@ def __init__(self, options, pluginmanager, start=True):
             self.options.cov_report = {}
         self.options.cov_source = _prepare_cov_source(self.options.cov_source)
 
+        if getattr(options, 'cov_contexts', None) == 'test':
+            self.test_contexts = True
+
         if is_dist and start:
             self.start(engine.DistMaster)
         elif start:
@@ -293,8 +311,12 @@ def pytest_runtest_setup(self, item):
             # test is run in another process than session, run
             # coverage manually
             embed.init()
+        if self.test_contexts:
+            self.switch_context(item, 'setup')
 
     def pytest_runtest_teardown(self, item):
+        if self.test_contexts:
+            self.switch_context(item, 'teardown')
         embed.cleanup()
 
     @compat.hookwrapper
@@ -305,8 +327,14 @@ def pytest_runtest_call(self, item):
             yield
             self.cov_controller.resume()
         else:
+            if self.test_contexts:
+                self.switch_context(item, 'run')
             yield
 
+    def switch_context(self, item, when):
+        context = "{item.nodeid}|{when}".format(item=item, when=when)
+        self.cov_controller.cov.switch_context(context)
+
 
 @pytest.fixture
 def no_cover():