asyn middleware

debug_toolbar/middleware.py

@@ -6,6 +6,7 @@
 import socket
 from functools import lru_cache
 
+from asgiref.sync import iscoroutinefunction, markcoroutinefunction
 from django.conf import settings
 from django.utils.module_loading import import_string
 
@@ -64,10 +65,30 @@ class DebugToolbarMiddleware:
     on outgoing response.
     """
 
+    sync_capable = True
+    async_capable = True
+
     def __init__(self, get_response):
+        if get_response is None:
+            raise ValueError("get_response must be provided.")
         self.get_response = get_response
+        # If get_response is a coroutine function, turns us into async mode so
+        # a thread is not consumed during a whole request.
+        self.async_mode = iscoroutinefunction(self.get_response)
+
+        if self.async_mode:
+            # Mark the class as async-capable, but do the actual switch inside
+            # __call__ to avoid swapping out dunder methods.
+            markcoroutinefunction(self)
+
+    def _enable_panel_instrumentation(self, panels):
+        for panel in panels:
+            panel.enable_instrumentation()
 
     def __call__(self, request):
+        if self.async_mode:
+            return self.__acall__(request)
+
         # Decide whether the toolbar is active for this request.
         show_toolbar = get_show_toolbar()
         if not show_toolbar(request) or DebugToolbar.is_toolbar_request(request):
@@ -123,6 +144,64 @@ def __call__(self, request):
                 response["Content-Length"] = len(response.content)
         return response
 
+    async def __acall__(self, request):
+        # Decide whether the toolbar is active for this request.
+        show_toolbar = get_show_toolbar()
+        if not show_toolbar(request) or DebugToolbar.is_toolbar_request(request):
+            response = await self.get_response(request)
+            return response
+
+        toolbar = DebugToolbar(request, self.get_response)
+
+        # Activate instrumentation ie. monkey-patch.
+        for panel in toolbar.enabled_panels:
+            panel.enable_instrumentation()
+        try:
+            # Run panels like Django middleware.
+            print(request)
+            response = await toolbar.process_request(request)
+        finally:
+            clear_stack_trace_caches()
+            # Deactivate instrumentation ie. monkey-unpatch. This must run
+            # regardless of the response. Keep 'return' clauses below.
+            for panel in reversed(toolbar.enabled_panels):
+                panel.disable_instrumentation()
+
+        # Generate the stats for all requests when the toolbar is being shown,
+        # but not necessarily inserted.
+        for panel in reversed(toolbar.enabled_panels):
+            panel.generate_stats(request, response)
+            panel.generate_server_timing(request, response)
+
+        # Always render the toolbar for the history panel, even if it is not
+        # included in the response.
+        rendered = toolbar.render_toolbar()
+
+        for header, value in self.get_headers(request, toolbar.enabled_panels).items():
+            response.headers[header] = value
+
+        # Check for responses where the toolbar can't be inserted.
+        content_encoding = response.get("Content-Encoding", "")
+        content_type = response.get("Content-Type", "").split(";")[0]
+        if (
+            getattr(response, "streaming", False)
+            or content_encoding != ""
+            or content_type not in _HTML_TYPES
+        ):
+            return response
+
+        # Insert the toolbar in the response.
+        content = response.content.decode(response.charset)
+        insert_before = dt_settings.get_config()["INSERT_BEFORE"]
+        pattern = re.escape(insert_before)
+        bits = re.split(pattern, content, flags=re.IGNORECASE)
+        if len(bits) > 1:
+            bits[-2] += rendered
+            response.content = insert_before.join(bits)
+            if "Content-Length" in response:
+                response["Content-Length"] = len(response.content)
+        return response
+
     @staticmethod
     def get_headers(request, panels):
         headers = {}