Fix several issues relating to access to source code inside zipfiles. Initial work by Alexander Belopolsky. See Misc/NEWS in this checkin for details.

8 files changed, 111 insertions, 48 deletions

diff --git a/Lib/bdb.py b/Lib/bdb.py
index 5288cc0..f29fa46 100644
--- a/Lib/bdb.py
+++ b/Lib/bdb.py
@@ -347,7 +347,7 @@ class Bdb:
             rv = frame.f_locals['__return__']
             s = s + '->'
             s = s + repr.repr(rv)
-        line = linecache.getline(filename, lineno)
+        line = linecache.getline(filename, lineno, frame.f_globals)
         if line: s = s + lprefix + line.strip()
         return s
 
@@ -589,7 +589,7 @@ class Tdb(Bdb):
         name = frame.f_code.co_name
         if not name: name = '???'
         fn = self.canonic(frame.f_code.co_filename)
-        line = linecache.getline(fn, frame.f_lineno)
+        line = linecache.getline(fn, frame.f_lineno, frame.f_globals)
         print '+++', fn, frame.f_lineno, name, ':', line.strip()
     def user_return(self, frame, retval):
         print '+++ return', retval
diff --git a/Lib/linecache.py b/Lib/linecache.py
index 4838625..48f7dda 100644
--- a/Lib/linecache.py
+++ b/Lib/linecache.py
@@ -88,21 +88,20 @@ def updatecache(filename, module_globals=None):
             get_source = getattr(loader, 'get_source', None)
 
             if name and get_source:
-                if basename.startswith(name.split('.')[-1]+'.'):
-                    try:
-                        data = get_source(name)
-                    except (ImportError, IOError):
-                        pass
-                    else:
-                        if data is None:
-                            # No luck, the PEP302 loader cannot find the source
-                            # for this module.
-                            return []
-                        cache[filename] = (
-                            len(data), None,
-                            [line+'\n' for line in data.splitlines()], fullname
-                        )
-                        return cache[filename][2]
+                try:
+                    data = get_source(name)
+                except (ImportError, IOError):
+                    pass
+                else:
+                    if data is None:
+                        # No luck, the PEP302 loader cannot find the source
+                        # for this module.
+                        return []
+                    cache[filename] = (
+                        len(data), None,
+                        [line+'\n' for line in data.splitlines()], fullname
+                    )
+                    return cache[filename][2]
 
         # Try looking through the module search path.
 
diff --git a/Lib/pdb.py b/Lib/pdb.py
index 4a080c7..3f76032 100755
--- a/Lib/pdb.py
+++ b/Lib/pdb.py
@@ -440,7 +440,7 @@ class Pdb(bdb.Bdb, cmd.Cmd):
         Return `lineno` if it is, 0 if not (e.g. a docstring, comment, blank
         line or EOF). Warning: testing is not comprehensive.
         """
-        line = linecache.getline(filename, lineno)
+        line = linecache.getline(filename, lineno, self.curframe.f_globals)
         if not line:
             print >>self.stdout, 'End of file'
             return 0
@@ -768,7 +768,7 @@ class Pdb(bdb.Bdb, cmd.Cmd):
         breaklist = self.get_file_breaks(filename)
         try:
             for lineno in range(first, last+1):
-                line = linecache.getline(filename, lineno)
+                line = linecache.getline(filename, lineno, self.curframe.f_globals)
                 if not line:
                     print >>self.stdout, '[EOF]'
                     break
diff --git a/Lib/runpy.py b/Lib/runpy.py
index f3c3890..fea6104 100755
--- a/Lib/runpy.py
+++ b/Lib/runpy.py
@@ -65,13 +65,14 @@ def _run_module_code(code, init_globals=None,
 
 # This helper is needed due to a missing component in the PEP 302
 # loader protocol (specifically, "get_filename" is non-standard)
+# Since we can't introduce new features in maintenance releases,
+# support was added to zipimporter under the name '_get_filename'
 def _get_filename(loader, mod_name):
-    try:
-        get_filename = loader.get_filename
-    except AttributeError:
-        return None
-    else:
-        return get_filename(mod_name)
+    for attr in ("get_filename", "_get_filename"):
+        meth = getattr(loader, attr, None)
+        if meth is not None:
+            return meth(mod_name)
+    return None
 
 # Helper to get the loader, code and filename for a module
 def _get_module_details(mod_name):
diff --git a/Lib/test/test_cmd_line_script.py b/Lib/test/test_cmd_line_script.py
index ba410ab..995d4a8 100644
--- a/Lib/test/test_cmd_line_script.py
+++ b/Lib/test/test_cmd_line_script.py
@@ -75,36 +75,66 @@ def _compile_test_script(script_name):
         compiled_name = script_name + 'o'
     return compiled_name
 
-def _make_test_zip(zip_dir, zip_basename, script_name):
+def _make_test_zip(zip_dir, zip_basename, script_name, name_in_zip=None):
     zip_filename = zip_basename+os.extsep+'zip'
     zip_name = os.path.join(zip_dir, zip_filename)
     zip_file = zipfile.ZipFile(zip_name, 'w')
-    zip_file.write(script_name, os.path.basename(script_name))
+    if name_in_zip is None:
+        name_in_zip = os.path.basename(script_name)
+    zip_file.write(script_name, name_in_zip)
     zip_file.close()
-    # if verbose:
+    #if verbose:
     #    zip_file = zipfile.ZipFile(zip_name, 'r')
     #    print 'Contents of %r:' % zip_name
     #    zip_file.printdir()
     #    zip_file.close()
-    return zip_name
+    return zip_name, os.path.join(zip_name, name_in_zip)
 
 def _make_test_pkg(pkg_dir):
     os.mkdir(pkg_dir)
     _make_test_script(pkg_dir, '__init__', '')
 
+def _make_test_zip_pkg(zip_dir, zip_basename, pkg_name, script_basename,
+                       source=test_source, depth=1):
+    init_name = _make_test_script(zip_dir, '__init__', '')
+    init_basename = os.path.basename(init_name)
+    script_name = _make_test_script(zip_dir, script_basename, source)
+    pkg_names = [os.sep.join([pkg_name]*i) for i in range(1, depth+1)]
+    script_name_in_zip = os.path.join(pkg_names[-1], os.path.basename(script_name))
+    zip_filename = zip_basename+os.extsep+'zip'
+    zip_name = os.path.join(zip_dir, zip_filename)
+    zip_file = zipfile.ZipFile(zip_name, 'w')
+    for name in pkg_names:
+        init_name_in_zip = os.path.join(name, init_basename)
+        zip_file.write(init_name, init_name_in_zip)
+    zip_file.write(script_name, script_name_in_zip)
+    zip_file.close()
+    os.unlink(init_name)
+    os.unlink(script_name)
+    #if verbose:
+    #    zip_file = zipfile.ZipFile(zip_name, 'r')
+    #    print 'Contents of %r:' % zip_name
+    #    zip_file.printdir()
+    #    zip_file.close()
+    return zip_name, os.path.join(zip_name, script_name_in_zip)
+
 # There's no easy way to pass the script directory in to get
 # -m to work (avoiding that is the whole point of making
 # directories and zipfiles executable!)
 # So we fake it for testing purposes with a custom launch script
 launch_source = """\
 import sys, os.path, runpy
-sys.path[0:0] = os.path.dirname(__file__)
+sys.path.insert(0, %s)
 runpy._run_module_as_main(%r)
 """
 
-def _make_launch_script(script_dir, script_basename, module_name):
-    return _make_test_script(script_dir, script_basename,
-                             launch_source % module_name)
+def _make_launch_script(script_dir, script_basename, module_name, path=None):
+    if path is None:
+        path = "os.path.dirname(__file__)"
+    else:
+        path = repr(path)
+    source = launch_source % (path, module_name)
+    return _make_test_script(script_dir, script_basename, source)
 
 class CmdLineTest(unittest.TestCase):
     def _check_script(self, script_name, expected_file,
@@ -155,15 +185,15 @@ class CmdLineTest(unittest.TestCase):
     def test_zipfile(self):
         with temp_dir() as script_dir:
             script_name = _make_test_script(script_dir, '__main__')
-            zip_name = _make_test_zip(script_dir, 'test_zip', script_name)
-            self._check_script(zip_name, None, zip_name, '')
+            zip_name, run_name = _make_test_zip(script_dir, 'test_zip', script_name)
+            self._check_script(zip_name, run_name, zip_name, '')
 
     def test_zipfile_compiled(self):
         with temp_dir() as script_dir:
             script_name = _make_test_script(script_dir, '__main__')
             compiled_name = _compile_test_script(script_name)
-            zip_name = _make_test_zip(script_dir, 'test_zip', compiled_name)
-            self._check_script(zip_name, None, zip_name, '')
+            zip_name, run_name = _make_test_zip(script_dir, 'test_zip', compiled_name)
+            self._check_script(zip_name, run_name, zip_name, '')
 
     def test_module_in_package(self):
         with temp_dir() as script_dir:
@@ -171,8 +201,19 @@ class CmdLineTest(unittest.TestCase):
             _make_test_pkg(pkg_dir)
             script_name = _make_test_script(pkg_dir, 'script')
             launch_name = _make_launch_script(script_dir, 'launch', 'test_pkg.script')
-            self._check_script(launch_name, script_name,
-                               script_name, 'test_pkg')
+            self._check_script(launch_name, script_name, script_name, 'test_pkg')
+
+    def test_module_in_package_in_zipfile(self):
+        with temp_dir() as script_dir:
+            zip_name, run_name = _make_test_zip_pkg(script_dir, 'test_zip', 'test_pkg', 'script')
+            launch_name = _make_launch_script(script_dir, 'launch', 'test_pkg.script', zip_name)
+            self._check_script(launch_name, run_name, run_name, 'test_pkg')
+
+    def test_module_in_subpackage_in_zipfile(self):
+        with temp_dir() as script_dir:
+            zip_name, run_name = _make_test_zip_pkg(script_dir, 'test_zip', 'test_pkg', 'script', depth=2)
+            launch_name = _make_launch_script(script_dir, 'launch', 'test_pkg.test_pkg.script', zip_name)
+            self._check_script(launch_name, run_name, run_name, 'test_pkg.test_pkg')
 
 
 def test_main():
diff --git a/Lib/test/test_doctest.py b/Lib/test/test_doctest.py
index 943fb82..a812868 100644
--- a/Lib/test/test_doctest.py
+++ b/Lib/test/test_doctest.py
@@ -6,6 +6,9 @@ from test import test_support
 import doctest
 import warnings
 
+# NOTE: There are some additional tests relating to interaction with
+#       zipimport in the test_zipimport_support test module.
+
 ######################################################################
 ## Sample Objects (used by test cases)
 ######################################################################
@@ -369,7 +372,7 @@ We'll simulate a __file__ attr that ends in pyc:
     >>> tests = finder.find(sample_func)
 
     >>> print tests  # doctest: +ELLIPSIS
-    [<DocTest sample_func from ...:13 (1 example)>]
+    [<DocTest sample_func from ...:16 (1 example)>]
 
 The exact name depends on how test_doctest was invoked, so allow for
 leading path components.
diff --git a/Lib/test/test_inspect.py b/Lib/test/test_inspect.py
index 32eeb57..300d143 100644
--- a/Lib/test/test_inspect.py
+++ b/Lib/test/test_inspect.py
@@ -16,6 +16,9 @@ from test import inspect_fodder2 as mod2
 # getclasstree, getargspec, getargvalues, formatargspec, formatargvalues,
 # currentframe, stack, trace, isdatadescriptor
 
+# NOTE: There are some additional tests relating to interaction with
+#       zipimport in the test_zipimport_support test module.
+
 modfile = mod.__file__
 if modfile.endswith(('c', 'o')):
     modfile = modfile[:-1]
diff --git a/Lib/test/test_zipimport.py b/Lib/test/test_zipimport.py
index 205853d..87869ae 100644
--- a/Lib/test/test_zipimport.py
+++ b/Lib/test/test_zipimport.py
@@ -214,16 +214,24 @@ class UncompressedZipImportTestCase(ImportHooksBaseTestCase):
             zi = zipimport.zipimporter(TEMP_ZIP)
             self.assertEquals(zi.archive, TEMP_ZIP)
             self.assertEquals(zi.is_package(TESTPACK), True)
-            zi.load_module(TESTPACK)
+            mod = zi.load_module(TESTPACK)
+            self.assertEquals(zi._get_filename(TESTPACK), mod.__file__)
 
             self.assertEquals(zi.is_package(packdir + '__init__'), False)
             self.assertEquals(zi.is_package(packdir + TESTPACK2), True)
             self.assertEquals(zi.is_package(packdir2 + TESTMOD), False)
 
-            mod_name = packdir2 + TESTMOD
-            mod = __import__(module_path_to_dotted_name(mod_name))
+            mod_path = packdir2 + TESTMOD
+            mod_name = module_path_to_dotted_name(mod_path)
+            pkg = __import__(mod_name)
+            mod = sys.modules[mod_name]
             self.assertEquals(zi.get_source(TESTPACK), None)
-            self.assertEquals(zi.get_source(mod_name), None)
+            self.assertEquals(zi.get_source(mod_path), None)
+            self.assertEquals(zi._get_filename(mod_path), mod.__file__)
+            # To pass in the module name instead of the path, we must use the right importer
+            loader = mod.__loader__
+            self.assertEquals(loader.get_source(mod_name), None)
+            self.assertEquals(loader._get_filename(mod_name), mod.__file__)
 
             # test prefix and archivepath members
             zi2 = zipimport.zipimporter(TEMP_ZIP + os.sep + TESTPACK)
@@ -251,15 +259,23 @@ class UncompressedZipImportTestCase(ImportHooksBaseTestCase):
             self.assertEquals(zi.archive, TEMP_ZIP)
             self.assertEquals(zi.prefix, packdir)
             self.assertEquals(zi.is_package(TESTPACK2), True)
-            zi.load_module(TESTPACK2)
+            mod = zi.load_module(TESTPACK2)
+            self.assertEquals(zi._get_filename(TESTPACK2), mod.__file__)
 
             self.assertEquals(zi.is_package(TESTPACK2 + os.sep + '__init__'), False)
             self.assertEquals(zi.is_package(TESTPACK2 + os.sep + TESTMOD), False)
 
-            mod_name = TESTPACK2 + os.sep + TESTMOD
-            mod = __import__(module_path_to_dotted_name(mod_name))
+            mod_path = TESTPACK2 + os.sep + TESTMOD
+            mod_name = module_path_to_dotted_name(mod_path)
+            pkg = __import__(mod_name)
+            mod = sys.modules[mod_name]
             self.assertEquals(zi.get_source(TESTPACK2), None)
-            self.assertEquals(zi.get_source(mod_name), None)
+            self.assertEquals(zi.get_source(mod_path), None)
+            self.assertEquals(zi._get_filename(mod_path), mod.__file__)
+            # To pass in the module name instead of the path, we must use the right importer
+            loader = mod.__loader__
+            self.assertEquals(loader.get_source(mod_name), None)
+            self.assertEquals(loader._get_filename(mod_name), mod.__file__)
         finally:
             z.close()
             os.remove(TEMP_ZIP)