Custom implementation of inspect.stack() which safely handles errors with findsource

1 files changed, 59 insertions, 3 deletions

diff --git a/debug_toolbar/utils/__init__.py b/debug_toolbar/utils/__init__.py
index c4dc160..0a78dac 100644
--- a/debug_toolbar/utils/__init__.py
+++ b/debug_toolbar/utils/__init__.py
@@ -1,6 +1,8 @@
+import inspect
 import os.path
 import django
 import SocketServer
+import sys
 
 from django.conf import settings
 from django.views.debug import linebreak_iter
@@ -21,7 +23,7 @@ def tidy_stacktrace(stack):
     1. Are part of Django (except contrib apps)
     2. Are part of SocketServer (used by Django's dev server)
     3. Are the last entry (which is part of our stacktracing code)
-    
+
     ``stack`` should be a list of frame tuples from ``inspect.stack()``
     """
     trace = []
@@ -84,9 +86,63 @@ def get_name_from_obj(obj):
         name = obj.__class__.__name__
     else:
         name = '<unknown>'
-    
+
     if hasattr(obj, '__module__'):
         module = obj.__module__
         name = '%s.%s' % (module, name)
 
-    return name
\ No newline at end of file
+    return name
+
+def getframeinfo(frame, context=1):
+    """
+    Get information about a frame or traceback object.
+
+    A tuple of five things is returned: the filename, the line number of
+    the current line, the function name, a list of lines of context from
+    the source code, and the index of the current line within that list.
+    The optional second argument specifies the number of lines of context
+    to return, which are centered around the current line.
+
+    This originally comes from ``inspect`` but is modified to handle issues
+    with ``findsource()``.
+    """
+    if inspect.istraceback(frame):
+        lineno = frame.tb_lineno
+        frame = frame.tb_frame
+    else:
+        lineno = frame.f_lineno
+    if not inspect.isframe(frame):
+        raise TypeError('arg is not a frame or traceback object')
+
+    filename = inspect.getsourcefile(frame) or inspect.getfile(frame)
+    if context > 0:
+        start = lineno - 1 - context//2
+        try:
+            lines, lnum = inspect.findsource(frame)
+        except (IOError, IndexError):
+            lines = index = None
+        else:
+            start = max(start, 1)
+            start = max(0, min(start, len(lines) - context))
+            lines = lines[start:start+context]
+            index = lineno - 1 - start
+    else:
+        lines = index = None
+
+    return inspect.Traceback(filename, lineno, frame.f_code.co_name, lines, index)
+
+def get_stack(context=1):
+    """
+    Get a list of records for a frame and all higher (calling) frames.
+
+    Each record contains a frame object, filename, line number, function
+    name, a list of lines of context, and index within the context.
+
+    Modified version of ``inspect.stack()`` which calls our own ``getframeinfo()``
+    """
+    frame = sys._getframe(1)
+    framelist = []
+    while frame:
+        framelist.append((frame,) + getframeinfo(frame, context))
+        frame = frame.f_back
+    return framelist
\ No newline at end of file