source_for_file helper with unit tests

3 files changed, 29 insertions, 4 deletions

diff --git a/coverage/control.py b/coverage/control.py
index 91600ae..376d3f5 100644
--- a/coverage/control.py
+++ b/coverage/control.py
@@ -27,7 +27,7 @@ from coverage.misc import file_be_gone, isolate_module
 from coverage.multiproc import patch_multiprocessing
 from coverage.plugin import FileReporter
 from coverage.plugin_support import Plugins
-from coverage.python import PythonFileReporter
+from coverage.python import PythonFileReporter, source_for_file
 from coverage.results import Analysis, Numbers
 from coverage.summary import SummaryReporter
 from coverage.xmlreport import XmlReporter
@@ -422,7 +422,7 @@ class Coverage(object):
         # co_filename value.
         dunder_file = frame.f_globals.get('__file__')
         if dunder_file:
-            filename = self._source_for_file(dunder_file)
+            filename = source_for_file(dunder_file)
             if original_filename and not original_filename.startswith('<'):
                 orig = os.path.basename(original_filename)
                 if orig != os.path.basename(filename):
diff --git a/coverage/python.py b/coverage/python.py
index ada21e6..5e142d2 100644
--- a/coverage/python.py
+++ b/coverage/python.py
@@ -91,7 +91,7 @@ def get_zip_bytes(filename):
     return None
 
 
-def _source_for_file(self, filename):
+def source_for_file(filename):
     """Return the source file for `filename`.
 
     Given a file name being traced, return the best guess as to the source
diff --git a/tests/test_python.py b/tests/test_python.py
index ee1e1f9..c4e333e 100644
--- a/tests/test_python.py
+++ b/tests/test_python.py
@@ -5,8 +5,9 @@
 
 import os
 import sys
+from coverage import env
 
-from coverage.python import get_zip_bytes
+from coverage.python import get_zip_bytes, source_for_file
 
 from tests.coveragetest import CoverageTest
 
@@ -28,3 +29,27 @@ class GetZipBytesTest(CoverageTest):
             self.assertIn('All OK', zip_text)
             # Run the code to see that we really got it encoded properly.
             __import__("encoded_"+encoding)
+
+def test_source_for_file(tmpdir):
+    path = tmpdir.join("a.py")
+    src = str(path)
+    assert src == source_for_file(src)
+    assert src == source_for_file(src + 'c')
+    assert src == source_for_file(src + 'o')
+    unknown = src + 'FOO'
+    assert unknown == source_for_file(unknown)
+    #
+    # On windows if a pyw exists, it is an acceptable source
+    #
+    windows = env.WINDOWS
+    env.WINDOWS = True
+    path_windows = tmpdir.ensure("a.pyw")
+    assert str(path_windows) == source_for_file(src + 'c')
+    #
+    # If both pyw and py exist, py is preferred
+    #
+    path.ensure(file=True)
+    assert src == source_for_file(src + 'c')
+    env.WINDOWS = windows
+    jython_src = "a"
+    assert jython_src + ".py" == source_for_file(jython_src + "$py.class")