Improve template wrapping

Fix the issue seen in scoutapp/scout_apm_python#231 by using improved decoration of the `render()` methods on the template classes using the famous `wrapt` module.

Also improve the speed of getting the stack depth by avoiding pulling the whole of `inspect.stack()`, which does some reformatting work, and instead working with raw frames instaed.

Author: adamchainz

setup.py

@@ -14,6 +14,7 @@
     packages=find_packages(),
     install_requires=[
         'django-debug-toolbar>=1.1,<2.0',
+        'wrapt',
     ],
     include_package_data=True,
     zip_safe=False,                 # because we're including static files

template_profiler_panel/__init__.py

@@ -0,0 +1 @@
+default_app_config = "template_profiler_panel.apps.TemplateProfilerPanelAppConfig"

template_profiler_panel/apps.py

@@ -0,0 +1,52 @@
+import inspect
+from time import time
+
+import wrapt
+from django.apps import AppConfig
+
+from template_profiler_panel.signals import template_rendered
+
+
+class TemplateProfilerPanelAppConfig(AppConfig):
+    name = "template_profiler_panel"
+    verbose_name = "Debug Toolbar Template Profiler Panel"
+
+    def ready(self):
+        self.monkey_patch_template_classes()
+
+    def monkey_patch_template_classes(self):
+        from django.template import Template as DjangoTemplate
+        template_classes = [DjangoTemplate]
+
+        try:
+            from jinja2 import Template as Jinja2Template
+        except ImportError:
+            pass
+        else:
+            template_classes.append(Jinja2Template)
+
+        @wrapt.decorator
+        def render_wrapper(wrapped, instance, args, kwargs):
+            start = time()
+            result = wrapped(*args, **kwargs)
+            end = time()
+
+            stack_depth = 1
+            current_frame = inspect.currentframe()
+            while True:
+                current_frame = current_frame.f_back
+                if current_frame is None:
+                    break
+                stack_depth += 1
+
+            template_rendered.send(
+                sender=instance.__class__,
+                instance=instance,
+                start=start,
+                end=end,
+                level=stack_depth,
+            )
+            return result
+
+        for template_class in template_classes:
+            template_class.render = render_wrapper(template_class.render)

template_profiler_panel/panels/template.py

@@ -1,57 +1,11 @@
-from time import time
 from collections import defaultdict
-from inspect import stack
 
-from django.dispatch import Signal
 from django.utils.translation import ugettext_lazy as _
 
 from debug_toolbar.panels import Panel
 from debug_toolbar.panels.sql.utils import contrasting_color_generator
 
-from django.template import Template as DjangoTemplate
-
-try:
-    jinja_import = True
-    from jinja2 import Template as JinjaTemplate
-except ImportError:
-    jinja_import = False
-
-template_rendered = Signal(
-    providing_args=['instance', 'start', 'end', 'level'])
-
-
-def template_render_wrapper_django(self, context):
-    t_start = time()
-    result = DjangoTemplate.tp_saved_render(self, context)
-    t_end = time()
-
-    template_rendered.send(
-        sender=DjangoTemplate, instance=self, start=t_start, end=t_end,
-        level=len(stack()))
-
-    return result
-
-
-def template_render_wrapper_jinja(self, context):
-    t_start = time()
-    result = JinjaTemplate.tp_saved_render(self, context)
-    t_end = time()
-
-    template_rendered.send(
-        sender=JinjaTemplate, instance=self, start=t_start, end=t_end,
-        level=len(stack()))
-
-    return result
-
-
-DjangoTemplate.engine = DjangoTemplate
-DjangoTemplate.tp_saved_render = DjangoTemplate.render
-DjangoTemplate.render = template_render_wrapper_django
-
-if jinja_import:
-    JinjaTemplate.engine = JinjaTemplate
-    JinjaTemplate.tp_saved_render = JinjaTemplate.render
-    JinjaTemplate.render = template_render_wrapper_jinja
+from template_profiler_panel.signals import template_rendered
 
 
 class TemplateProfilerPanel(Panel):

template_profiler_panel/signals.py

@@ -0,0 +1,3 @@
+from django.dispatch import Signal
+
+template_rendered = Signal(providing_args=['instance', 'start', 'end', 'level'])