Issue #17911: traceback module overhaul

Provide a way to seed the linecache for a PEP-302 module without actually
loading the code.

Provide a new object API for traceback, including the ability to not lookup
lines at all until the traceback is actually rendered, without any trace of the
original objects being kept alive.

Author: rbtcollins

linecache2/__init__.py

@@ -5,6 +5,7 @@
 that name.
 """
 
+import functools
 import io
 import sys
 import os
@@ -22,7 +23,9 @@ def getline(filename, lineno, module_globals=None):
 
 # The cache
 
-cache = {} # The cache
+# The cache. Maps filenames to either a thunk which will provide source code,
+# or a tuple (size, mtime, lines, fullname) once loaded.
+cache = {}
 
 
 def clearcache():
@@ -37,6 +40,9 @@ def getlines(filename, module_globals=None):
     Update the cache if it doesn't contain an entry for this file already."""
 
     if filename in cache:
+        entry = cache[filename]
+        if len(entry) == 1:
+            return updatecache(filename, module_globals)
         return cache[filename][2]
     else:
         return updatecache(filename, module_globals)
@@ -55,7 +61,11 @@ def checkcache(filename=None):
             return
 
     for filename in filenames:
-        size, mtime, lines, fullname = cache[filename]
+        entry = cache[filename]
+        if len(entry) == 1:
+            # lazy cache entry, leave it lazy.
+            continue
+        size, mtime, lines, fullname = entry
         if mtime is None:
             continue   # no-op for files loaded via a __loader__
         try:
@@ -73,7 +83,8 @@ def updatecache(filename, module_globals=None):
     and return an empty list."""
 
     if filename in cache:
-        del cache[filename]
+        if len(cache[filename]) != 1:
+            del cache[filename]
     if not filename or (filename.startswith('<') and filename.endswith('>')):
         return []
 
@@ -83,27 +94,23 @@ def updatecache(filename, module_globals=None):
     except OSError:
         basename = filename
 
-        # Try for a __loader__, if available
-        if module_globals and '__loader__' in module_globals:
-            name = module_globals.get('__name__')
-            loader = module_globals['__loader__']
-            get_source = getattr(loader, 'get_source', None)
-
-            if name and get_source:
-                try:
-                    data = get_source(name)
-                except (ImportError, OSError):
-                    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]
+        # Realise a lazy loader based lookup if there is one
+        # otherwise try to lookup right now.
+        if lazycache(filename, module_globals):
+            try:
+                data = cache[filename][0]()
+            except (ImportError, OSError):
+                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, which is only useful
         # when handling a relative filename.
@@ -134,6 +141,40 @@ def updatecache(filename, module_globals=None):
     cache[filename] = size, mtime, lines, fullname
     return lines
 
+
+def lazycache(filename, module_globals):
+    """Seed the cache for filename with module_globals.
+
+    The module loader will be asked for the source only when getlines is
+    called, not immediately.
+
+    If there is an entry in the cache already, it is not altered.
+
+    :return: True if a lazy load is registered in the cache,
+        otherwise False. To register such a load a module loader with a
+        get_source method must be found, the filename must be a cachable
+        filename, and the filename must not be already cached.
+    """
+    if filename in cache:
+        if len(cache[filename]) == 1:
+            return True
+        else:
+            return False
+    if not filename or (filename.startswith('<') and filename.endswith('>')):
+        return False
+    # Try for a __loader__, if available
+    if module_globals and '__loader__' in module_globals:
+        name = module_globals.get('__name__')
+        loader = module_globals['__loader__']
+        get_source = getattr(loader, 'get_source', None)
+
+        if name and get_source:
+            get_lines = functools.partial(get_source, name)
+            cache[filename] = (get_lines,)
+            return True
+    return False
+
+
 #### ---- avoiding having a tokenize2 backport for now ----
 from codecs import lookup, BOM_UTF8
 import re

linecache2/tests/test_linecache.py

@@ -10,6 +10,7 @@
 FILENAME = os.__file__
 if FILENAME.endswith('.pyc'):
     FILENAME = FILENAME[:-1]
+NONEXISTENT_FILENAME = FILENAME + '.missing'
 INVALID_NAME = '!@$)(!@#_1'
 EMPTY = ''
 TESTS = 'inspect_fodder inspect_fodder2 mapping_tests'
@@ -137,3 +138,47 @@ def test_checkcache(self):
             for index, line in enumerate(source):
                 self.assertEqual(line, getline(source_name, index + 1))
                 source_list.append(line)
+ 
+    def test_lazycache_no_globals(self):
+        lines = linecache.getlines(FILENAME)
+        linecache.clearcache()
+        self.assertEqual(False, linecache.lazycache(FILENAME, None))
+        self.assertEqual(lines, linecache.getlines(FILENAME))
+
+    @unittest.skipIf("__loader__" not in globals(), "Modules not PEP302 by default")
+    def test_lazycache_smoke(self):
+        lines = linecache.getlines(NONEXISTENT_FILENAME, globals())
+        linecache.clearcache()
+        self.assertEqual(
+            True, linecache.lazycache(NONEXISTENT_FILENAME, globals()))
+        self.assertEqual(1, len(linecache.cache[NONEXISTENT_FILENAME]))
+        # Note here that we're looking up a non existant filename with no
+        # globals: this would error if the lazy value wasn't resolved.
+        self.assertEqual(lines, linecache.getlines(NONEXISTENT_FILENAME))
+
+    def test_lazycache_provide_after_failed_lookup(self):
+        linecache.clearcache()
+        lines = linecache.getlines(NONEXISTENT_FILENAME, globals())
+        linecache.clearcache()
+        linecache.getlines(NONEXISTENT_FILENAME)
+        linecache.lazycache(NONEXISTENT_FILENAME, globals())
+        self.assertEqual(lines, linecache.updatecache(NONEXISTENT_FILENAME))
+
+    def test_lazycache_check(self):
+        linecache.clearcache()
+        linecache.lazycache(NONEXISTENT_FILENAME, globals())
+        linecache.checkcache()
+
+    def test_lazycache_bad_filename(self):
+        linecache.clearcache()
+        self.assertEqual(False, linecache.lazycache('', globals()))
+        self.assertEqual(False, linecache.lazycache('<foo>', globals()))
+
+    @unittest.skipIf("__loader__" not in globals(), "Modules not PEP302 by default")
+    def test_lazycache_already_cached(self):
+        linecache.clearcache()
+        lines = linecache.getlines(NONEXISTENT_FILENAME, globals())
+        self.assertEqual(
+            False,
+            linecache.lazycache(NONEXISTENT_FILENAME, globals()))
+        self.assertEqual(4, len(linecache.cache[NONEXISTENT_FILENAME]))