Use some strange self.os stuff to hold references to modules so we'll have them available when called from atexit(). Fixes issue #71, I think.

2 files changed, 17 insertions, 7 deletions

diff --git a/coverage/control.py b/coverage/control.py
index 7ab046cb..b7dfa0a8 100644
--- a/coverage/control.py
+++ b/coverage/control.py
@@ -29,7 +29,6 @@ class coverage(object):
         cov.html_report(directory='covhtml')
 
     """
-
     def __init__(self, data_file=None, data_suffix=None, cover_pylib=None,
                 auto_data=False, timid=None, branch=None, config_file=True,
                 source=None, omit=None, include=None):
@@ -164,6 +163,10 @@ class coverage(object):
         # Only _harvest_data once per measurement cycle.
         self._harvested = False
 
+        self.socket = socket
+        self.os = os
+        self.random = random
+
     def _canonical_dir(self, f):
         """Return the canonical directory of the file `f`."""
         return os.path.split(self.file_locator.canonical_filename(f))[0]
@@ -185,6 +188,9 @@ class coverage(object):
         should not.
 
         """
+        if os is None:
+            return False
+
         if filename.startswith('<'):
             # Lots of non-file execution is represented with artificial
             # filenames like "<string>", "<doctest readme.txt[0]>", or
@@ -366,13 +372,13 @@ 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 is True:
             # 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.socket.gethostname(), self.os.getpid(), self.random.randint(0, 99999)
                     )
 
         self._harvest_data()
diff --git a/coverage/data.py b/coverage/data.py
index 1b883750..fe076bef 100644
--- a/coverage/data.py
+++ b/coverage/data.py
@@ -56,6 +56,10 @@ class CoverageData(object):
         #
         self.arcs = {}
 
+        self.os = os
+        self.sorted = sorted
+        self.pickle = pickle
+
     def usefile(self, use_file=True):
         """Set whether or not to use a disk file for data."""
         self.use_file = use_file
@@ -93,13 +97,13 @@ class CoverageData(object):
     def line_data(self):
         """Return the map from filenames to lists of line numbers executed."""
         return dict(
-            [(f, sorted(lmap.keys())) for f, lmap in self.lines.items()]
+            [(f, self.sorted(lmap.keys())) for f, lmap in self.lines.items()]
             )
 
     def arc_data(self):
         """Return the map from filenames to lists of line number pairs."""
         return dict(
-            [(f, sorted(amap.keys())) for f, amap in self.arcs.items()]
+            [(f, self.sorted(amap.keys())) for f, amap in self.arcs.items()]
             )
 
     def write_file(self, filename):
@@ -119,7 +123,7 @@ class CoverageData(object):
         # Write the pickle to the file.
         fdata = open(filename, 'wb')
         try:
-            pickle.dump(data, fdata, 2)
+            self.pickle.dump(data, fdata, 2)
         finally:
             fdata.close()
 
@@ -229,7 +233,7 @@ class CoverageData(object):
         if fullpath:
             filename_fn = lambda f: f
         else:
-            filename_fn = os.path.basename
+            filename_fn = self.os.path.basename
         for filename, lines in self.lines.items():
             summ[filename_fn(filename)] = len(lines)
         return summ