1 files changed, 37 insertions, 21 deletions

diff --git a/coverage/codeunit.py b/coverage/codeunit.py
index 28fa0551..01708957 100644
--- a/coverage/codeunit.py
+++ b/coverage/codeunit.py
@@ -6,22 +6,27 @@ from coverage.backward import string_class, StringIO
 from coverage.misc import CoverageException
 
 
-def code_unit_factory(morfs, file_locator, omit_prefixes=None):
+def code_unit_factory(
+    morfs, file_locator, omit_prefixes=None, include_prefixes=None
+    ):
     """Construct a list of CodeUnits from polymorphic inputs.
-    
+
     `morfs` is a module or a filename, or a list of same.
+
     `file_locator` is a FileLocator that can help resolve filenames.
-    `omit_prefixes` is a list of prefixes.  CodeUnits that match those prefixes
-    will be omitted from the list.
-    
+
+    `include_prefixes` is a list of prefixes. Only CodeUnits that match those
+    prefixes will be included in the list. `omit_prefixes` is a list of
+    prefixes to omit from the list.
+
     Returns a list of CodeUnit objects.
-    
+
     """
 
     # Be sure we have a list.
     if not isinstance(morfs, (list, tuple)):
         morfs = [morfs]
-    
+
     # On Windows, the shell doesn't expand wildcards.  Do it here.
     globbed = []
     for morf in morfs:
@@ -32,7 +37,18 @@ def code_unit_factory(morfs, file_locator, omit_prefixes=None):
     morfs = globbed
 
     code_units = [CodeUnit(morf, file_locator) for morf in morfs]
-    
+
+    if include_prefixes:
+        assert not isinstance(include_prefixes, string_class) # common mistake
+        prefixes = [file_locator.abs_file(p) for p in include_prefixes]
+        filtered = []
+        for cu in code_units:
+            for prefix in prefixes:
+                if cu.filename.startswith(prefix):
+                    filtered.append(cu)
+                    break
+        code_units = filtered
+
     if omit_prefixes:
         code_units = omit_filter(omit_prefixes, code_units)
 
@@ -57,13 +73,13 @@ def omit_filter(omit_prefixes, code_units):
 
 class CodeUnit(object):
     """Code unit: a filename or module.
-    
+
     Instance attributes:
-    
+
     `name` is a human-readable name for this code unit.
     `filename` is the os path from which we can read the source.
     `relative` is a boolean.
-    
+
     """
 
     def __init__(self, morf, file_locator):
@@ -98,40 +114,40 @@ class CodeUnit(object):
 
     # Annoying comparison operators. Py3k wants __lt__ etc, and Py2k needs all
     # of them defined.
-    
+
     def __lt__(self, other):
         return self.name < other.name
-    
+
     def __le__(self, other):
         return self.name <= other.name
 
     def __eq__(self, other):
         return self.name == other.name
-    
+
     def __ne__(self, other):
         return self.name != other.name
 
     def __gt__(self, other):
         return self.name > other.name
-    
+
     def __ge__(self, other):
         return self.name >= other.name
 
     def flat_rootname(self):
         """A base for a flat filename to correspond to this code unit.
-        
+
         Useful for writing files about the code where you want all the files in
         the same directory, but need to differentiate same-named files from
         different directories.
-        
+
         For example, the file a/b/c.py might return 'a_b_c'
-        
+
         """
         if self.modname:
             return self.modname.replace('.', '_')
         else:
-            root = os.path.splitdrive(os.path.splitext(self.name)[0])[1]
-            return root.replace('\\', '_').replace('/', '_')
+            root = os.path.splitdrive(self.name)[1]
+            return root.replace('\\', '_').replace('/', '_').replace('.', '_')
 
     def source_file(self):
         """Return an open file for reading the source of the code unit."""
@@ -143,7 +159,7 @@ class CodeUnit(object):
         source = self.file_locator.get_zip_data(self.filename)
         if source is not None:
             return StringIO(source)
-            
+
         # Couldn't find source.
         raise CoverageException(
             "No source for code %r." % self.filename