chore: clean up impl.

Author: taeold

src/functions_framework/aio/__init__.py

@@ -16,12 +16,9 @@
 import contextvars
 import functools
 import inspect
-import json
 import logging
 import logging.config
 import os
-import re
-import sys
 import traceback
 
 from typing import Any, Awaitable, Callable, Dict, Tuple, Union

src/functions_framework/execution_id.py

@@ -136,8 +136,17 @@ async def __call__(self, scope, receive, send):
         await self.app(scope, receive, send)  # pragma: no cover
 
 
-# Sets execution id and span id for the request
 def set_execution_context(request, enable_id_logging=False):
+    """Decorator for Flask/WSGI handlers that sets execution context.
+
+    Takes request object at decoration time (Flask pattern where request is available
+    via thread-local context when decorator is applied).
+
+    Usage:
+        @set_execution_context(request, enable_id_logging=True)
+        def view_func(path):
+            ...
+    """
     if enable_id_logging:
         stdout_redirect = contextlib.redirect_stdout(
             LoggingHandlerAddExecutionId(sys.stdout)
@@ -165,6 +174,16 @@ def wrapper(*args, **kwargs):
 
 
 def set_execution_context_async(enable_id_logging=False):
+    """Decorator for ASGI/async handlers that sets execution context.
+
+    Unlike set_execution_context which takes request at decoration time (Flask pattern),
+    this expects the decorated function to receive request as its first parameter (ASGI pattern).
+
+    Usage:
+        @set_execution_context_async(enable_id_logging=True)
+        async def handler(request, *args, **kwargs):
+            ...
+    """
     if enable_id_logging:
         stdout_redirect = contextlib.redirect_stdout(
             LoggingHandlerAddExecutionId(sys.stdout)

tests/test_execution_id_async.py

@@ -15,7 +15,6 @@
 import json
 import pathlib
 import re
-import sys
 
 from functools import partial
 from unittest.mock import Mock, patch
@@ -297,14 +296,15 @@ async def test_maintains_execution_id_for_concurrent_requests(monkeypatch, capsy
 
 def test_async_decorator_with_sync_function():
     """Test that the async decorator handles sync functions properly."""
-    from functions_framework.execution_id import set_execution_context_async
 
     # Create a sync function
     def sync_func(request):
         return {"status": "ok"}
 
-    # Apply the async decorator
-    wrapped = set_execution_context_async(enable_id_logging=False)(sync_func)
+    # Apply the decorator
+    wrapped = execution_id.set_execution_context_async(enable_id_logging=False)(
+        sync_func
+    )
 
     # Create mock request
     request = Mock()