Calculate the pid suffix for data files at the end of the process so that programs calling os.fork will collect data from both child and parent. Fixes issue #56.

5 files changed, 48 insertions, 35 deletions

diff --git a/CHANGES.txt b/CHANGES.txt
index e705ef36..92ac8828 100644
--- a/CHANGES.txt
+++ b/CHANGES.txt
@@ -15,14 +15,19 @@ Next version
   line, which is highlighted on arrival.  Added a link back to the index page
   at the bottom of each HTML page.
 
+- Programs that call ``os.fork`` will properly collect data from both the child
+  and parent processes.  Use ``coverage run -p`` to get two data files that can
+  be combined with ``coverage combine``.  Fixes `issue 56`_.
+
 - Source files can have more than one dot in them (foo.test.py), and will be
   treated properly while reporting.  Fixes `issue 46`_.
 
 - Source files with DOS line endings are now properly tokenized for syntax
-  coloring on non-DOS machines.  Fixes `issue  53`_.
+  coloring on non-DOS machines.  Fixes `issue 53`_.
 
 .. _issue 46: http://bitbucket.org/ned/coveragepy/issue/46
 .. _issue 53: http://bitbucket.org/ned/coveragepy/issue/53
+.. _issue 56: http://bitbucket.org/ned/coveragepy/issue/56
 
 
 Version 3.3.1, 6 March 2010
diff --git a/coverage/control.py b/coverage/control.py
index 1456178a..a3287ae6 100644
--- a/coverage/control.py
+++ b/coverage/control.py
@@ -96,17 +96,22 @@ class coverage(object):
             branch=self.config.branch
             )
 
-        # Create the data file.
+        # Suffixes are a bit tricky.  We want to use the data suffix only when
+        # collecting data, not when combining data.  So we save it as
+        # `self.run_suffix` now, and promote it to `self.data_suffix` if we
+        # find that we are collecting data later.
         if data_suffix or self.config.parallel:
             if not isinstance(data_suffix, string_class):
                 # if data_suffix=True, use .machinename.pid.random
-                data_suffix = "%s.%s.%06d" % (
-                    socket.gethostname(), os.getpid(), random.randint(0, 99999)
-                    )
+                data_suffix = True
         else:
             data_suffix = None
+        self.data_suffix = None
         self.run_suffix = data_suffix
 
+        # Create the data file.  We do this at construction time so that the
+        # data file will be written into the directory where the process
+        # started rather than wherever the process eventually chdir'd to.
         self.data = CoverageData(
             basename=self.config.data_file,
             collector="coverage v%s" % __version__
@@ -192,13 +197,9 @@ class coverage(object):
     def start(self):
         """Start measuring code coverage."""
         if self.run_suffix:
-            # If the .coveragerc file specifies parallel=True, then we need to
-            # remake the data file for collection, with a suffix.
-            from coverage import __version__
-            self.data = CoverageData(
-                basename=self.config.data_file, suffix=self.run_suffix,
-                collector="coverage v%s" % __version__
-                )
+            # Calling start() means we're running code, so use the run_suffix
+            # as the data_suffix when we eventually save the data.
+            self.data_suffix = self.run_suffix
         if self.auto_data:
             self.load()
             # Save coverage data when Python exits.
@@ -249,8 +250,18 @@ class coverage(object):
 
     def save(self):
         """Save the collected coverage data to the data file."""
+        data_suffix = self.data_suffix
+        if data_suffix and not isinstance(data_suffix, string_class):
+            # If data_suffix was a simple true value, then make a suffix with
+            # plenty of distinguishing information.  We do this here in
+            # `save()` at the last minute so that the pid will be correct even
+            # if the process forks.
+            data_suffix = "%s.%s.%06d" % (
+                    socket.gethostname(), os.getpid(), random.randint(0, 99999)
+                    )
+
         self._harvest_data()
-        self.data.write()
+        self.data.write(suffix=data_suffix)
 
     def combine(self):
         """Combine together a number of similarly-named coverage data files.
diff --git a/coverage/data.py b/coverage/data.py
index 9359af12..1b883750 100644
--- a/coverage/data.py
+++ b/coverage/data.py
@@ -21,22 +21,11 @@ class CoverageData(object):
 
     """
 
-    # Name of the data file (unless environment variable is set).
-    filename_default = ".coverage"
-
-    # Environment variable naming the data file.
-    filename_env = "COVERAGE_FILE"
-
-    def __init__(self, basename=None, suffix=None, collector=None):
+    def __init__(self, basename=None, collector=None):
         """Create a CoverageData.
 
         `basename` is the name of the file to use for storing data.
 
-        `suffix` is a suffix to append to the base file name. This can be used
-        for multiple or parallel execution, so that many coverage data files
-        can exist simultaneously.  A dot will be used to join the base name and
-        the suffix.
-
         `collector` is a string describing the coverage measurement software.
 
         """
@@ -47,8 +36,6 @@ class CoverageData(object):
         # Construct the filename that will be used for data file storage, if we
         # ever do any file storage.
         self.filename = basename or ".coverage"
-        if suffix:
-            self.filename += "." + suffix
         self.filename = os.path.abspath(self.filename)
 
         # A map from canonical Python source file name to a dictionary in
@@ -80,10 +67,20 @@ class CoverageData(object):
         else:
             self.lines, self.arcs = {}, {}
 
-    def write(self):
-        """Write the collected coverage data to a file."""
+    def write(self, suffix=None):
+        """Write the collected coverage data to a file.
+
+        `suffix` is a suffix to append to the base file name. This can be used
+        for multiple or parallel execution, so that many coverage data files
+        can exist simultaneously.  A dot will be used to join the base name and
+        the suffix.
+
+        """
         if self.use_file:
-            self.write_file(self.filename)
+            filename = self.filename
+            if suffix:
+                filename += "." + suffix
+            self.write_file(filename)
 
     def erase(self):
         """Erase the data, both in this object, and from its file storage."""
diff --git a/test/test_data.py b/test/test_data.py
index 4f784253..83a5b8ae 100644
--- a/test/test_data.py
+++ b/test/test_data.py
@@ -55,13 +55,13 @@ class DataTest(CoverageTest):
         self.assert_summary(covdata2, SUMMARY_1)
 
     def test_combining(self):
-        covdata1 = CoverageData(suffix='1')
+        covdata1 = CoverageData()
         covdata1.add_line_data(DATA_1)
-        covdata1.write()
+        covdata1.write(suffix='1')
 
-        covdata2 = CoverageData(suffix='2')
+        covdata2 = CoverageData()
         covdata2.add_line_data(DATA_2)
-        covdata2.write()
+        covdata2.write(suffix='2')
 
         covdata3 = CoverageData()
         covdata3.combine_parallel_data()
diff --git a/test/test_process.py b/test/test_process.py
index 501075df..1f8a9884 100644
--- a/test/test_process.py
+++ b/test/test_process.py
@@ -227,7 +227,7 @@ class ProcessTest(CoverageTest):
             self.assertEqual(out, 'Child!\n')
             self.assertFalse(os.path.exists(".coverage"))
 
-            # After running the forking program, there should be two 
+            # After running the forking program, there should be two
             # .coverage.machine.123 files.
             self.assertEqual(self.number_of_data_files(), 2)