fix: when checking source existence for remapping, zipfiles are ok

2 files changed, 43 insertions, 15 deletions

diff --git a/coverage/files.py b/coverage/files.py
index 63e9afb2..5f222419 100644
--- a/coverage/files.py
+++ b/coverage/files.py
@@ -153,6 +153,35 @@ def abs_file(path):
     return actual_path(os.path.abspath(os.path.realpath(path)))
 
 
+def zip_location(filename):
+    """Split a filename into a zipfile / inner name pair.
+
+    Only return a pair if the zipfile exists.  No check is made if the inner
+    name is in the zipfile.
+
+    """
+    for ext in ['.zip', '.egg', '.pex']:
+        zipbase, extension, inner = filename.partition(ext + sep(filename))
+        if extension:
+            zipfile = zipbase + ext
+            if os.path.exists(zipfile):
+                return zipfile, inner
+    return None
+
+
+def source_exists(path):
+    """Determine if a source file path exists."""
+    if os.path.exists(path):
+        return True
+
+    if zip_location(path):
+        # If zip_location returns anything, then it's a zipfile that
+        # exists. That's good enough for us.
+        return True
+
+    return False
+
+
 def python_reported_file(filename):
     """Return the string as Python would describe this file name."""
     if env.PYBEHAVIOR.report_absolute_files:
@@ -408,7 +437,7 @@ class PathAliases:
         result = result.rstrip(r"\/") + result_sep
         self.aliases.append((original_pattern, regex, result))
 
-    def map(self, path, exists=os.path.exists):
+    def map(self, path, exists=source_exists):
         """Map `path` through the aliases.
 
         `path` is checked against all of the patterns.  The first pattern to
diff --git a/coverage/python.py b/coverage/python.py
index c8b8e774..b3232085 100644
--- a/coverage/python.py
+++ b/coverage/python.py
@@ -9,7 +9,7 @@ import zipimport
 
 from coverage import env
 from coverage.exceptions import CoverageException, NoSource
-from coverage.files import canonical_filename, relative_filename
+from coverage.files import canonical_filename, relative_filename, zip_location
 from coverage.misc import contract, expensive, isolate_module, join_regex
 from coverage.parser import PythonParser
 from coverage.phystokens import source_token_lines, source_encoding
@@ -79,19 +79,18 @@ def get_zip_bytes(filename):
     an empty string if the file is empty.
 
     """
-    markers = ['.zip'+os.sep, '.egg'+os.sep, '.pex'+os.sep]
-    for marker in markers:
-        if marker in filename:
-            parts = filename.split(marker)
-            try:
-                zi = zipimport.zipimporter(parts[0]+marker[:-1])
-            except zipimport.ZipImportError:
-                continue
-            try:
-                data = zi.get_data(parts[1])
-            except OSError:
-                continue
-            return data
+    zipfile_inner = zip_location(filename)
+    if zipfile_inner is not None:
+        zipfile, inner = zipfile_inner
+        try:
+            zi = zipimport.zipimporter(zipfile)
+        except zipimport.ZipImportError:
+            return None
+        try:
+            data = zi.get_data(inner)
+        except OSError:
+            return None
+        return data
     return None