quart: support both legacy and non-legacy

Author: bitterpanda63

aikido_zen/sources/quart.py

@@ -1,17 +1,34 @@
+import inspect
 from aikido_zen.context import get_current_context
 from .functions.asgi_middleware import InternalASGIMiddleware
 from ..helpers.get_argument import get_argument
-from ..sinks import on_import, patch_function, before_async
+from ..sinks import on_import, patch_function, before_async, after
 
 
-async def _asgi_app(func, instance, args, kwargs):
+async def _call_coroutine(func, instance, args, kwargs):
     scope = get_argument(args, kwargs, 0, "scope")
     receive = get_argument(args, kwargs, 1, "receive")
     send = get_argument(args, kwargs, 2, "send")
 
     await InternalASGIMiddleware(func, "quart")(scope, receive, send)
 
 
+@after
+def _call(func, instance, args, kwargs, return_value):
+    """
+    Legacy ASGI v2.0
+    func: application(scope)
+    return_value: coroutine application_instance(receive, send)
+    """
+    scope = get_argument(args, kwargs, 0, "scope")
+
+    async def application_instance(receive, send):
+        await InternalASGIMiddleware(return_value, "quart")(scope, receive, send)
+
+    # Modify return_value
+    return_value = application_instance
+
+
 @before_async
 async def _handle_request_before(func, instance, args, kwargs):
     context = get_current_context()
@@ -37,9 +54,19 @@ async def _handle_request_before(func, instance, args, kwargs):
 @on_import("quart.app", "quart")
 def patch(m):
     """
-    patching module quart.app
-    - patches Quart.asgi_app (handles internal asgi middleware)
+    We patch Quart.__call__ instead of asgi_app, because asgi_app itself can be wrapped multiple times
+    And we want to be the first middleware to run.
+    - patches Quart.__call__ (handles internal asgi middleware)
     - patches Quart.handle_request (Stores body/cookies)
     """
-    patch_function(m, "Quart.asgi_app", _asgi_app)
+
+    if inspect.iscoroutine(m.Quart.__call__):
+        # coroutine application(scope, receive, send)
+        patch_function(m, "Quart.__call__", _call_coroutine)
+    else:
+        # Legacy ASGI v2.0
+        # https://asgi.readthedocs.io/en/latest/specs/main.html#legacy-applications
+        # application(scope): coroutine application_instance(receive, send)
+        patch_function(m, "Quart.__call__", _call)
+
     patch_function(m, "Quart.handle_request", _handle_request_before)