feat: better support use of aio decorator

Author: taeold

src/functions_framework/__init__.py

@@ -327,6 +327,16 @@ def crash_handler(e):
 
 
 def create_app(target=None, source=None, signature_type=None):
+    """Create an app for the function.
+
+    Args:
+        target: The name of the target function to invoke
+        source: The source file containing the function
+        signature_type: The signature type of the function
+
+    Returns:
+        A Flask WSGI app or Starlette ASGI app depending on function decorators
+    """
     target = _function_registry.get_function_target(target)
     source = _function_registry.get_function_source(source)
 
@@ -370,6 +380,7 @@ def handle_none(rv):
         setup_logging()
 
     _app.wsgi_app = execution_id.WsgiMiddleware(_app.wsgi_app)
+
     # Execute the module, within the application context
     with _app.app_context():
         try:
@@ -394,6 +405,23 @@ def function(*_args, **_kwargs):
                 # command fails.
                 raise e from None
 
+    use_asgi = target in _function_registry.ASGI_FUNCTIONS
+    if use_asgi:
+        # This function needs ASGI, delegate to create_asgi_app
+        # Note: @aio decorators only register functions in ASGI_FUNCTIONS when the
+        # module is imported. We can't know if a function uses @aio until after
+        # we load the module.
+        #
+        # To avoid loading modules twice, we always create a Flask app first, load the
+        # module within its context, then check if ASGI is needed. This results in an
+        # unused Flask app for ASGI functions, but we accept this memory overhead as a
+        # trade-off.
+        from functions_framework.aio import create_asgi_app_from_module
+
+        return create_asgi_app_from_module(
+            target, source, signature_type, source_module, spec
+        )
+
     # Get the configured function signature type
     signature_type = _function_registry.get_func_signature_type(target, signature_type)
 

src/functions_framework/_cli.py

@@ -39,14 +39,10 @@
     help="Use ASGI server for function execution",
 )
 def _cli(target, source, signature_type, host, port, debug, asgi):
-    if not asgi and target in _function_registry.ASGI_FUNCTIONS:
-        asgi = True
-
-    if asgi:  # pragma: no cover
+    if asgi:
         from functions_framework.aio import create_asgi_app
 
         app = create_asgi_app(target, source, signature_type)
     else:
         app = create_app(target, source, signature_type)
-
     create_server(app, debug).run(host, port)

src/functions_framework/aio/__init__.py

@@ -215,6 +215,29 @@ async def __call__(self, scope, receive, send):
             # Don't re-raise to prevent starlette from printing traceback again
 
 
+def create_asgi_app_from_module(target, source, signature_type, source_module, spec):
+    """Create an ASGI application from an already-loaded module.
+
+    Args:
+        target: The name of the target function to invoke
+        source: The source file containing the function
+        signature_type: The signature type of the function
+        source_module: The already-loaded module
+        spec: The module spec
+
+    Returns:
+        A Starlette ASGI application instance
+    """
+    enable_id_logging = _enable_execution_id_logging()
+    if enable_id_logging:
+        _configure_app_execution_id_logging()
+
+    function = _function_registry.get_user_function(source, source_module, target)
+    signature_type = _function_registry.get_func_signature_type(target, signature_type)
+
+    return _create_asgi_app_with_function(function, signature_type, enable_id_logging)
+
+
 def create_asgi_app(target=None, source=None, signature_type=None):
     """Create an ASGI application for the function.
 
@@ -245,6 +268,11 @@ def create_asgi_app(target=None, source=None, signature_type=None):
     function = _function_registry.get_user_function(source, source_module, target)
     signature_type = _function_registry.get_func_signature_type(target, signature_type)
 
+    return _create_asgi_app_with_function(function, signature_type, enable_id_logging)
+
+
+def _create_asgi_app_with_function(function, signature_type, enable_id_logging):
+    """Create an ASGI app with the given function and signature type."""
     is_async = inspect.iscoroutinefunction(function)
     routes = []
     if signature_type == _function_registry.HTTP_SIGNATURE_TYPE: