Properly report .pyw files. #261.

1 files changed, 40 insertions, 8 deletions

diff --git a/coverage/results.py b/coverage/results.py
index 77ff2a2d..2d13e818 100644
--- a/coverage/results.py
+++ b/coverage/results.py
@@ -15,16 +15,10 @@ class Analysis(object):
         self.code_unit = code_unit
 
         self.filename = self.code_unit.filename
-        ext = os.path.splitext(self.filename)[1]
-        source = None
-        if ext == '.py':
-            if not os.path.exists(self.filename):
-                source = self.coverage.file_locator.get_zip_data(self.filename)
-                if not source:
-                    raise NoSource("No source for code: '%s'" % self.filename)
+        actual_filename, source = self.find_source(self.filename)
 
         self.parser = CodeParser(
-            text=source, filename=self.filename,
+            text=source, filename=actual_filename,
             exclude=self.coverage._exclude_regex('exclude')
             )
         self.statements, self.excluded = self.parser.parse_source()
@@ -59,6 +53,44 @@ class Analysis(object):
             n_missing_branches=n_missing_branches,
             )
 
+    def find_source(self, filename):
+        """Find the source for `filename`.
+
+        Returns two values: the actual filename, and the source.
+
+        The source returned depends on which of these cases holds:
+
+            * The filename seems to be a non-source file: returns None
+
+            * The filename is a source file, and actually exists: returns None.
+
+            * The filename is a source file, and is in a zip file or egg:
+              returns the source.
+
+            * The filename is a source file, but couldn't be found: raises
+              `NoSource`.
+
+        """
+        source = None
+
+        base, ext = os.path.splitext(filename)
+        TRY_EXTS = {
+            '.py':  ['.py', '.pyw'],
+            '.pyw': ['.pyw'],
+        }
+        try_exts = TRY_EXTS.get(ext)
+        if not try_exts:
+            return filename, None
+
+        for try_ext in try_exts:
+            try_filename = base + try_ext
+            if os.path.exists(try_filename):
+                return try_filename, None
+            source = self.coverage.file_locator.get_zip_data(try_filename)
+            if source:
+                return try_filename, source
+        raise NoSource("No source for code: '%s'" % filename)
+
     def missing_formatted(self):
         """The missing line numbers, formatted nicely.