Wildly experimental multiprocessing support. Covers most of #117.

6 files changed, 110 insertions, 3 deletions

diff --git a/AUTHORS.txt b/AUTHORS.txt
index f2d561e5..49181897 100644
--- a/AUTHORS.txt
+++ b/AUTHORS.txt
@@ -44,6 +44,7 @@ Brandon Rhodes
 Adi Roiban
 Greg Rogers
 Chris Rose
+Eduardo Schettino
 George Song
 Anthony Sottile
 David Stanek
diff --git a/CHANGES.txt b/CHANGES.txt
index b0004b5c..013820c5 100644
--- a/CHANGES.txt
+++ b/CHANGES.txt
@@ -6,9 +6,15 @@ Change history for Coverage.py
 Latest
 ------
 
-- A new warning is possible, if a desired file isn't measure because it was
+- Wildly experimental: support for measuring processes started by the
+  multiprocessing module.  To use, set ``--concurrency=multiprocessing``,
+  either on the command line or in the .coveragerc file. Thanks, Eduardo
+  Schettino. (`issue 117`_).
+
+- A new warning is possible, if a desired file isn't measured because it was
   imported before coverage was started (`issue 353`_).
 
+.. _issue 117: https://bitbucket.org/ned/coveragepy/issue/117/enable-coverage-measurement-of-code-run-by
 .. _issue 353: https://bitbucket.org/ned/coveragepy/issue/353/40a3-introduces-an-unexpected-third-case
 
 
diff --git a/coverage/cmdline.py b/coverage/cmdline.py
index ff3de4ca..4ef08f1d 100644
--- a/coverage/cmdline.py
+++ b/coverage/cmdline.py
@@ -24,7 +24,9 @@ class Opts(object):
         '', '--branch', action='store_true',
         help="Measure branch coverage in addition to statement coverage."
         )
-    CONCURRENCY_CHOICES = ["thread", "gevent", "greenlet", "eventlet"]
+    CONCURRENCY_CHOICES = [
+        "thread", "gevent", "greenlet", "eventlet", "multiprocessing",
+    ]
     concurrency = optparse.make_option(
         '', '--concurrency', action='store', metavar="LIB",
         choices=CONCURRENCY_CHOICES,
diff --git a/coverage/control.py b/coverage/control.py
index 319f56dc..f422d7c0 100644
--- a/coverage/control.py
+++ b/coverage/control.py
@@ -21,6 +21,7 @@ from coverage.files import ModuleMatcher
 from coverage.html import HtmlReporter
 from coverage.misc import CoverageException, bool_or_none, join_regex
 from coverage.misc import file_be_gone, overrides
+from coverage.monkey import patch_multiprocessing
 from coverage.plugin import CoveragePlugin, FileReporter
 from coverage.python import PythonCodeUnit
 from coverage.results import Analysis, Numbers
@@ -219,13 +220,18 @@ class Coverage(object):
         self.omit = prep_patterns(self.config.omit)
         self.include = prep_patterns(self.config.include)
 
+        concurrency = self.config.concurrency
+        if concurrency == "multiprocessing":
+            patch_multiprocessing()
+            concurrency = None
+
         self.collector = Collector(
             should_trace=self._should_trace,
             check_include=self._check_include_omit_etc,
             timid=self.config.timid,
             branch=self.config.branch,
             warn=self._warn,
-            concurrency=self.config.concurrency,
+            concurrency=concurrency,
             )
 
         # Suffixes are a bit tricky.  We want to use the data suffix only when
@@ -544,6 +550,8 @@ class Coverage(object):
     def _warn(self, msg):
         """Use `msg` as a warning."""
         self._warnings.append(msg)
+        if self.debug.should("pid"):
+            msg = "[%d] %s" % (os.getpid(), msg)
         sys.stderr.write("Coverage.py warning: %s\n" % msg)
 
     def use_cache(self, usecache):
diff --git a/coverage/monkey.py b/coverage/monkey.py
new file mode 100644
index 00000000..42f185ea
--- /dev/null
+++ b/coverage/monkey.py
@@ -0,0 +1,45 @@
+"""Monkey-patching to make coverage work right in some cases."""
+
+import multiprocessing
+import multiprocessing.process
+import sys
+
+# An attribute that will be set on modules to indicate that they have been
+# monkey-patched.
+MARKER = "_coverage$patched"
+
+
+def patch_multiprocessing():
+    """Monkey-patch the multiprocessing module.
+
+    This enables coverage measurement of processes started by multiprocessing.
+    This is wildly experimental!
+
+    """
+    if hasattr(multiprocessing, MARKER):
+        return
+
+    if sys.version_info >= (3, 4):
+        klass = multiprocessing.process.BaseProcess
+    else:
+        klass = multiprocessing.Process
+
+    original_bootstrap = klass._bootstrap
+
+    class ProcessWithCoverage(klass):
+        def _bootstrap(self):
+            from coverage import Coverage
+            cov = Coverage(data_suffix=True)
+            cov.start()
+            try:
+                return original_bootstrap(self)
+            finally:
+                cov.stop()
+                cov.save()
+
+    if sys.version_info >= (3, 4):
+        klass._bootstrap = ProcessWithCoverage._bootstrap
+    else:
+        multiprocessing.Process = ProcessWithCoverage
+
+    setattr(multiprocessing, MARKER, 1)
diff --git a/tests/test_concurrency.py b/tests/test_concurrency.py
index 77b8c0ec..928f9404 100644
--- a/tests/test_concurrency.py
+++ b/tests/test_concurrency.py
@@ -220,6 +220,51 @@ class ConcurrencyTest(CoverageTest):
         self.try_some_code(BUG_330, "eventlet", eventlet, "0\n")
 
 
+class MultiprocessingTest(CoverageTest):
+    """Test support of the multiprocessing module."""
+
+    def test_multiprocessing(self):
+        self.make_file("multi.py", """\
+            import multiprocessing
+            import os
+            import time
+
+            def func(x):
+                # Need to pause, or the tasks go too quick, and some processes
+                # in the pool don't get any work, and then don't record data.
+                time.sleep(0.01)
+                if x % 2:
+                    return os.getpid(), x*x
+                else:
+                    return os.getpid(), x*x
+
+            if __name__ == "__main__":
+                pool = multiprocessing.Pool(3)
+                inputs = range(20)
+                outputs = pool.imap_unordered(func, inputs)
+                pids = set()
+                total = 0
+                for pid, sq in outputs:
+                    pids.add(pid)
+                    total += sq
+                print("%d pids, total = %d" % (len(pids), total))
+                pool.close()
+                pool.join()
+            """)
+
+        out = self.run_command(
+            "coverage run --concurrency=multiprocessing multi.py"
+        )
+        os.system("cp .cov* /tmp")
+        total = sum(x*x for x in range(20))
+        self.assertEqual(out.rstrip(), "3 pids, total = %d" % total)
+
+        self.run_command("coverage combine")
+        out = self.run_command("coverage report -m")
+        last_line = self.squeezed_lines(out)[-1]
+        self.assertEqual(last_line, "multi.py 20 0 100%")
+
+
 def print_simple_annotation(code, linenos):
     """Print the lines in `code` with X for each line number in `linenos`."""
     for lineno, line in enumerate(code.splitlines(), start=1):