feat(utils,websocket): create reuseable websocket utils and add transform function for playground to transform websocket route info

Author: huynguyengl99

chanx/playground/utils.py

@@ -1,9 +1,8 @@
 """
-Utility functions for WebSocket route discovery and inspection.
+WebSocket playground utilities.
 
-This module provides tools to dynamically discover WebSocket routes
-in a Django Channels application and generate example messages for
-WebSocket consumer endpoints.
+This module provides specialized utilities for the WebSocket playground,
+transforming route information into a format suitable for display and interaction.
 """
 
 import inspect
@@ -17,14 +16,14 @@
     get_type_hints,
 )
 
-from channels.routing import URLRouter, get_default_application
 from django.http import HttpRequest
 
 from polyfactory.factories.pydantic_factory import ModelFactory
 from pydantic import BaseModel
 
 from chanx.messages.base import BaseIncomingMessage
 from chanx.utils.logging import logger
+from chanx.utils.websocket import RouteInfo, get_websocket_routes, transform_routes
 
 
 class MessageExample(TypedDict):
@@ -51,185 +50,51 @@ class WebSocketRoute(TypedDict, total=False):
     message_examples: list[MessageExample]
 
 
-def get_websocket_routes(request: HttpRequest | None = None) -> list[WebSocketRoute]:
+def get_playground_websocket_routes(
+    request: HttpRequest | None = None,
+) -> list[WebSocketRoute]:
     """
-    Extract all WebSocket routes from the ASGI application.
+    Get WebSocket routes formatted for the playground.
 
-    This function traverses the Django Channels routing configuration
-    to discover all available WebSocket endpoints, their paths, and
-    generates example messages for each endpoint based on the
-    message schema defined in the consumer.
+    Uses the core WebSocket route discovery mechanism and transforms the routes
+    into a format suitable for the playground UI, including example messages.
 
     Args:
         request: The HTTP request object, used to determine the current domain.
                If None, defaults to localhost:8000.
 
     Returns:
-        A list of dictionaries containing route information including URL,
-        description, and example messages.
+        A list of WebSocketRoute objects with UI-friendly information.
     """
-    endpoints: list[WebSocketRoute] = []
+    # Get raw routes from the core utility
+    raw_routes = get_websocket_routes(request)
 
-    # Determine the WebSocket base URL based on the request
-    ws_base_url: str = _get_websocket_base_url(request)
+    # Transform routes into the format needed for the playground
+    return transform_routes(raw_routes, _transform_route_for_playground)
 
-    # Extract the WebSocket protocol handler from the ASGI application
-    application = get_default_application()
-    if hasattr(application, "application_mapping"):
-        # If it's a ProtocolTypeRouter
-        ws_app: Any = application.application_mapping.get("websocket")
 
-        # Extract routes
-        if ws_app:
-            _traverse_middleware(ws_app, "", endpoints, ws_base_url)
-
-    return endpoints
-
-
-def _get_websocket_base_url(request: HttpRequest | None) -> str:
-    """
-    Determine the WebSocket base URL based on the request.
-
-    Constructs a WebSocket URL (ws:// or wss://) based on the
-    domain in the request object.
-
-    Args:
-        request: The HTTP request object.
-
-    Returns:
-        The WebSocket base URL (ws:// or wss:// followed by domain).
-    """
-    if request is None:
-        return "ws://localhost:8000"
-
-    # Get the current domain from the request
-    domain: str = request.get_host()
-
-    # Determine if we should use secure WebSockets (wss://) based on the request
-    is_secure: bool = request.is_secure()
-    protocol: str = "wss://" if is_secure else "ws://"
-
-    return f"{protocol}{domain}"
-
-
-def _traverse_middleware(
-    app: Any, prefix: str, endpoints: list[WebSocketRoute], ws_base_url: str
-) -> None:
-    """
-    Traverse through middleware layers to find the URLRouter.
-
-    Recursively explores the middleware stack to find URLRouter instances
-    and extract route information from them. It uses the fact that
-    middleware typically stores the inner app as self.inner.
-
-    Args:
-        app: The current application or middleware to traverse.
-        prefix: URL prefix accumulated so far.
-        endpoints: List to store discovered endpoints.
-        ws_base_url: Base URL for WebSocket connections.
+def _transform_route_for_playground(route: RouteInfo) -> WebSocketRoute:
     """
-    # Skip if app is None
-    if app is None:
-        return
-
-    # If it's a URLRouter, extract routes
-    if isinstance(app, URLRouter):
-        _extract_routes_from_router(app, prefix, endpoints, ws_base_url)
-        return
-
-    # Try to access the inner application (standard middleware pattern)
-    inner_app: Any | None = getattr(app, "inner", None)
+    Transform a raw route into a playground-friendly format.
 
-    # If inner isn't found, try other common attributes that might hold the next app
-    if inner_app is None:
-        for attr_name in ["app", "application"]:
-            inner_app = getattr(app, attr_name, None)
-            if inner_app is not None:
-                break
-
-    # If we found an inner app, continue traversal
-    if inner_app is not None:
-        _traverse_middleware(inner_app, prefix, endpoints, ws_base_url)
-
-
-def _extract_routes_from_router(
-    router: URLRouter, prefix: str, endpoints: list[WebSocketRoute], ws_base_url: str
-) -> None:
-    """
-    Extract routes from a URLRouter object.
-
-    Processes each route in the router, extracting path patterns and
-    handler information, and recursively traversing nested routers.
-
-    Args:
-        router: The router to extract routes from.
-        prefix: URL prefix accumulated so far.
-        endpoints: List to store discovered endpoints.
-        ws_base_url: Base URL for WebSocket connections.
-    """
-    for route in router.routes:
-        try:
-            # Get the pattern string
-            pattern: str = _get_pattern_string(route)
-
-            # Build the full path
-            full_path: str = f"{prefix}{pattern}"
-
-            # Get the handler
-            handler: Any = route.callback
-
-            # If it's another router, recurse into it
-            if isinstance(handler, URLRouter):
-                _extract_routes_from_router(handler, full_path, endpoints, ws_base_url)
-            else:
-                # For consumers, get info with message examples
-                endpoint_info: WebSocketRoute = _get_handler_info(
-                    handler, full_path, ws_base_url
-                )
-                endpoints.append(endpoint_info)
-        except AttributeError as e:
-            # More specific error for attribute issues
-            logger.exception(
-                f"AttributeError while parsing route: {ws_base_url}/{prefix}. Error: {str(e)}"
-            )
-        except Exception as e:
-            # For other unexpected errors
-            logger.exception(
-                f"Error parsing route: {ws_base_url}/{prefix}. Error: {str(e)}"
-            )
-
-
-def _get_pattern_string(route: Any) -> str:
-    """
-    Extract pattern string from a route object.
-
-    Handles different route pattern implementations to extract
-    the URL pattern string.
+    This function extracts metadata from a WebSocket consumer and formats it
+    for display in the playground UI, including generating example messages.
 
     Args:
-        route: The route object to extract pattern from.
+        route: The RouteInfo dataclass instance
 
     Returns:
-        The cleaned URL pattern string.
+        A WebSocketRoute with UI-friendly information
     """
-    if hasattr(route, "pattern"):
-        # For URLRoute
-        if hasattr(route.pattern, "pattern"):
-            pattern: str = route.pattern.pattern
-        else:
-            # For RoutePattern
-            pattern = str(route.pattern)
-    else:
-        pattern = str(route)
-
-    # Clean up the pattern string
-    pattern = pattern.replace("^", "").replace("$", "")
-    return pattern
+    # Get handler info with examples for the playground
+    return _get_handler_info(
+        handler=route.handler, path=route.path, ws_base_url=route.base_url
+    )
 
 
 def _get_handler_info(handler: Any, path: str, ws_base_url: str) -> WebSocketRoute:
     """
-    Extract information about a route handler.
+    Extract information about a route handler for the playground.
 
     Extracts metadata from a WebSocket consumer including its name,
     description, and message schema, and generates example messages.
@@ -240,8 +105,7 @@ def _get_handler_info(handler: Any, path: str, ws_base_url: str) -> WebSocketRou
         ws_base_url: Base URL for WebSocket connections.
 
     Returns:
-        Information about the handler including name, URL, description,
-        and example messages.
+        Information about the handler formatted for the playground.
     """
     # Default values
     name: str = getattr(handler, "__name__", "Unknown")
@@ -283,7 +147,7 @@ def _create_example(msg_type: type[BaseModel]) -> MessageExample:
     Returns:
         A formatted example with name, description, and sample data.
     """
-    description = inspect.getdoc(msg_type) or f"Example of {msg_type.__name__}"
+    description: str = inspect.getdoc(msg_type) or f"Example of {msg_type.__name__}"
 
     # Create the example using the factory
     factory = ModelFactory.create_factory(model=msg_type)

chanx/playground/views.py

@@ -9,7 +9,7 @@
 
 from chanx.utils.logging import logger
 
-from .utils import WebSocketRoute, get_websocket_routes
+from .utils import WebSocketRoute, get_playground_websocket_routes
 
 
 class WebSocketPlaygroundView(TemplateView):
@@ -57,8 +57,10 @@ class WebSocketInfoView(APIView):
 
     def get(self, request: HttpRequest) -> Response:
         try:
-            # Get available WebSocket endpoints using the utility function
-            available_endpoints: list[WebSocketRoute] = get_websocket_routes(request)
+            # Get available WebSocket endpoints using the new playground utility function
+            available_endpoints: list[WebSocketRoute] = get_playground_websocket_routes(
+                request
+            )
 
             # Use the list serializer directly
             serializer = WebSocketRouteListSerializer(available_endpoints)