feature: add-websocket-support

Author: Jerzy Góra

README.md

@@ -34,6 +34,7 @@ pip install fastapi-controllers
 ```python
 import uvicorn
 from fastapi import FastAPI, Response, status
+from fastapi.websockets import WebSocket
 
 from fastapi_controllers import Controller, get
 
@@ -43,6 +44,13 @@ class ExampleController(Controller):
     async def get_example(self) -> Response:
         return Response(status_code=status.HTTP_200_OK)
 
+    @websocket("/ws")
+    async def ws_example(websocket: WebSocket) -> None:
+        await websocket.accept()
+        while True:
+            data = await websocket.receive_text()
+            await websocket.send_text(f"Received: {data}")
+
 
 if __name__ == "__main__":
     app = FastAPI()
@@ -54,12 +62,12 @@ FastAPI's `APIRouter` is created and populated with API routes by the `Controlle
 
 ## Seamless integration
 
-The router-related parameters as well as those of HTTP request-specific decorators are expected to be the same as those used by `fastapi.APIRouter` and `fastapi.APIRouter.<request_method>`. Validation of the provided parameters is performed during initialization via the `inspect` module. This ensures compatibility with the FastAPI framework and prevents the introduction of a new, unnecessary naming convention.
+The router-related parameters as well as those of HTTP request-specific and websocket decorators are expected to be the same as those used by `fastapi.APIRouter`, `fastapi.APIRouter.<request_method>` and `fastapi.APIRouter.websocket`. Validation of the provided parameters is performed during initialization via the `inspect` module. This ensures compatibility with the FastAPI framework and prevents the introduction of a new, unnecessary naming convention.
 
-### Supported HTTP request methods
+### Available decorators
 
 ```python
-from fastapi_controllers import delete, get, head, options, patch, post, put, trace
+from fastapi_controllers import delete, get, head, options, patch, post, put, trace, websocket
 ```
 
 ## Use class variables to customize your APIRouter

fastapi_controllers/__init__.py

@@ -1,5 +1,5 @@
 from fastapi_controllers.controllers import Controller
-from fastapi_controllers.routing import delete, get, head, options, patch, post, put, trace
+from fastapi_controllers.routing import delete, get, head, options, patch, post, put, trace, websocket
 
 __all__ = [
     "Controller",
@@ -11,4 +11,5 @@
     "post",
     "put",
     "trace",
+    "websocket",
 ]

fastapi_controllers/controllers.py

@@ -1,10 +1,10 @@
 import inspect
 from enum import Enum
-from types import SimpleNamespace
 from typing import Any, Dict, List, Optional, Sequence, Union
 
 from fastapi import APIRouter, params
 
+from fastapi_controllers.definitions import HTTPRouteDefinition, RouteData, WebsocketRouteDefinition
 from fastapi_controllers.helpers import _replace_signature, _validate_against_apirouter_signature
 
 
@@ -33,12 +33,21 @@ def create_router(cls) -> APIRouter:
         router = APIRouter(**cls.__router_params__)  # type: ignore
         for _, func in inspect.getmembers(cls, predicate=inspect.isfunction):
             _replace_signature(cls, func)
-            api_route_data: Optional[SimpleNamespace] = getattr(func, "__api_route_data__", None)
-            if api_route_data:
-                router.add_api_route(
-                    api_route_data.args[0],
-                    func,
-                    *api_route_data.args[1:],
-                    **api_route_data.kwargs,
-                )
+            route_data: Optional[RouteData] = getattr(func, "__route_data__", None)
+            if route_data:
+                if isinstance(route_data.route_definition, HTTPRouteDefinition):
+                    router.add_api_route(
+                        route_data.route_args[0],
+                        func,
+                        *route_data.route_args[1:],
+                        methods=[route_data.route_definition.request_method],
+                        **route_data.route_kwargs,
+                    )
+                if isinstance(route_data.route_definition, WebsocketRouteDefinition):
+                    router.add_api_websocket_route(
+                        route_data.route_args[0],
+                        func,
+                        *route_data.route_args[1:],
+                        **route_data.route_kwargs,
+                    )
         return router

fastapi_controllers/definitions.py

@@ -0,0 +1,49 @@
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any, Dict, Tuple
+
+
+class HTTPRequestMethod(str, Enum):
+    DELETE = "DELETE"
+    GET = "GET"
+    HEAD = "HEAD"
+    OPTIONS = "OPTIONS"
+    PATCH = "PATCH"
+    POST = "POST"
+    PUT = "PUT"
+    TRACE = "TRACE"
+
+
+class RouteDefinition:
+    def __init__(self, *, binds: str) -> None:
+        self.binds = binds
+
+
+class HTTPRouteDefinition(RouteDefinition):
+    def __init__(self, *, binds: str, request_method: HTTPRequestMethod) -> None:
+        super().__init__(binds=binds)
+        self.request_method = request_method
+
+
+class WebsocketRouteDefinition(RouteDefinition):
+    ...
+
+
+@dataclass
+class RouteData:
+    route_definition: RouteDefinition
+    route_args: Tuple[Any, ...]
+    route_kwargs: Dict[str, Any]
+
+
+@dataclass
+class Route:
+    delete = HTTPRouteDefinition(binds="delete", request_method=HTTPRequestMethod.DELETE)
+    get = HTTPRouteDefinition(binds="get", request_method=HTTPRequestMethod.GET)
+    head = HTTPRouteDefinition(binds="head", request_method=HTTPRequestMethod.HEAD)
+    options = HTTPRouteDefinition(binds="options", request_method=HTTPRequestMethod.OPTIONS)
+    patch = HTTPRouteDefinition(binds="patch", request_method=HTTPRequestMethod.PATCH)
+    post = HTTPRouteDefinition(binds="post", request_method=HTTPRequestMethod.POST)
+    put = HTTPRouteDefinition(binds="put", request_method=HTTPRequestMethod.PUT)
+    trace = HTTPRouteDefinition(binds="trace", request_method=HTTPRequestMethod.TRACE)
+    websocket = WebsocketRouteDefinition(binds="websocket")

fastapi_controllers/routing.py

@@ -1,68 +1,77 @@
-from enum import Enum
-from types import SimpleNamespace
-from typing import Any, Callable
+from typing import Any, Callable, Dict, Tuple
 
+from fastapi_controllers.definitions import Route, RouteData, RouteDefinition
 from fastapi_controllers.helpers import _validate_against_apirouter_signature
 
 
-class _HTTPRequestMethod(str, Enum):
-    DELETE = "DELETE"
-    GET = "GET"
-    HEAD = "HEAD"
-    OPTIONS = "OPTIONS"
-    PATCH = "PATCH"
-    POST = "POST"
-    PUT = "PUT"
-    TRACE = "TRACE"
-
-
 class _RouteDecorator:
-    method: str
+    _route_definition: RouteDefinition
 
-    def __init_subclass__(cls, method: _HTTPRequestMethod) -> None:
+    def __init_subclass__(cls, route_definition: RouteDefinition) -> None:
         super().__init_subclass__()
-        cls.method = method
+        cls._route_definition = route_definition
 
     def __init__(self, *args: Any, **kwargs: Any) -> None:
-        self.args = args
-        if self.args and callable(self.args[0]):
-            raise TypeError("You must provide a path for the route.")
-        self.kwargs = kwargs
-        _validate_against_apirouter_signature(self.method.lower(), args=args, kwargs=kwargs)
+        self.route_args = args
+        self.route_kwargs = kwargs
+        _validate_against_apirouter_signature(self._route_definition.binds, args=args, kwargs=kwargs)
 
     def __call__(self, func: Callable[..., Any]) -> Callable[..., Any]:
-        self.kwargs["methods"] = [self.method]
-        func.__api_route_data__ = SimpleNamespace(args=self.args, kwargs=self.kwargs)  # type: ignore
+        func.__route_data__ = RouteData(  # type: ignore
+            route_definition=self._route_definition,
+            route_args=self.route_args,
+            route_kwargs=self.route_kwargs,
+        )
         return func
 
+    @property
+    def route_args(self) -> Tuple[Any, ...]:
+        return self._route_args
+
+    @route_args.setter
+    def route_args(self, value: Tuple[Any, ...]) -> None:
+        self._route_args = value
+
+    @property
+    def route_kwargs(self) -> Dict[str, Any]:
+        return self._route_kwargs
+
+    @route_kwargs.setter
+    def route_kwargs(self, value: Dict[str, Any]) -> None:
+        self._route_kwargs = value
+
+
+class delete(_RouteDecorator, route_definition=Route.delete):
+    ...
+
 
-class delete(_RouteDecorator, method=_HTTPRequestMethod.DELETE):
+class get(_RouteDecorator, route_definition=Route.get):
     ...
 
 
-class get(_RouteDecorator, method=_HTTPRequestMethod.GET):
+class head(_RouteDecorator, route_definition=Route.head):
     ...
 
 
-class head(_RouteDecorator, method=_HTTPRequestMethod.HEAD):
+class options(_RouteDecorator, route_definition=Route.options):
     ...
 
 
-class options(_RouteDecorator, method=_HTTPRequestMethod.OPTIONS):
+class patch(_RouteDecorator, route_definition=Route.patch):
     ...
 
 
-class patch(_RouteDecorator, method=_HTTPRequestMethod.PATCH):
+class post(_RouteDecorator, route_definition=Route.post):
     ...
 
 
-class post(_RouteDecorator, method=_HTTPRequestMethod.POST):
+class put(_RouteDecorator, route_definition=Route.put):
     ...
 
 
-class put(_RouteDecorator, method=_HTTPRequestMethod.PUT):
+class trace(_RouteDecorator, route_definition=Route.trace):
     ...
 
 
-class trace(_RouteDecorator, method=_HTTPRequestMethod.TRACE):
+class websocket(_RouteDecorator, route_definition=Route.websocket):
     ...

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "fastapi-controllers"
-version = "0.1.1"
+version = "0.2.0"
 description = "Simple Controller implementation for FastAPI"
 authors = ["Jerzy Góra <[email protected]>"]
 license = "MIT"

tests/functional/conftest.py

@@ -4,8 +4,9 @@
 import pytest
 from fastapi import Depends, FastAPI, Response, status
 from fastapi.testclient import TestClient
+from fastapi.websockets import WebSocket
 
-from fastapi_controllers import Controller, delete, get, head, options, patch, post, put, trace
+from fastapi_controllers import Controller, delete, get, head, options, patch, post, put, trace, websocket
 
 
 def sync_dependency() -> str:
@@ -97,6 +98,12 @@ async def test_put(self) -> Response:
     async def test_trace(self) -> Response:
         return Response(status_code=status.HTTP_200_OK)
 
+    @websocket("/ws")
+    async def websocket(websocket: WebSocket) -> None:  # type: ignore
+        await websocket.accept()
+        await websocket.send_json({"msg": "Hello WebSocket"})
+        await websocket.close()
+
 
 @pytest.fixture
 def sync_test_client() -> TestClient:

tests/functional/test_async.py

@@ -1,21 +1,28 @@
 import pytest
 from fastapi import status
 from fastapi.testclient import TestClient
+from starlette.websockets import WebSocketDisconnect
 
-from fastapi_controllers.routing import _HTTPRequestMethod
+from fastapi_controllers.definitions import HTTPRequestMethod
 
 
 def describe_test_controller_async() -> None:
-    @pytest.mark.parametrize("http_request_method", _HTTPRequestMethod.__members__.values())
+    @pytest.mark.parametrize("http_request_method", HTTPRequestMethod.__members__.values())
     def it_responds_to_http_methods(
         async_test_client: TestClient,
-        http_request_method: _HTTPRequestMethod,
+        http_request_method: HTTPRequestMethod,
     ) -> None:
         request_func = getattr(async_test_client, http_request_method.lower(), None)
         if request_func:
             response = request_func("/test-async")
             assert response.status_code == status.HTTP_200_OK
 
+    def it_supports_websockets(async_test_client: TestClient) -> None:
+        with pytest.raises(WebSocketDisconnect):
+            with async_test_client.websocket_connect("/ws") as websocket:
+                data = websocket.receive_json()
+                assert data == {"msg": "Hello WebSocket"}
+
     def it_resolves_async_dependencies(async_test_client: TestClient) -> None:
         response = async_test_client.get("/test-async")
         assert response.status_code == status.HTTP_200_OK

tests/functional/test_sync.py

@@ -2,14 +2,14 @@
 from fastapi import status
 from fastapi.testclient import TestClient
 
-from fastapi_controllers.routing import _HTTPRequestMethod
+from fastapi_controllers.definitions import HTTPRequestMethod
 
 
 def describe_test_controller_sync() -> None:
-    @pytest.mark.parametrize("http_request_method", _HTTPRequestMethod.__members__.values())
+    @pytest.mark.parametrize("http_request_method", HTTPRequestMethod.__members__.values())
     def it_responds_to_http_methods(
         sync_test_client: TestClient,
-        http_request_method: _HTTPRequestMethod,
+        http_request_method: HTTPRequestMethod,
     ) -> None:
         request_func = getattr(sync_test_client, http_request_method.lower(), None)
         if request_func:

tests/unit/test_controllers.py

@@ -5,7 +5,7 @@
 from pytest_mock import MockerFixture
 
 from fastapi_controllers.controllers import Controller
-from fastapi_controllers.routing import get
+from fastapi_controllers.routing import get, websocket
 
 
 @pytest.fixture(autouse=True)
@@ -80,6 +80,10 @@ class FakeController(Controller):
                 def fake_method(self) -> None:
                     ...
 
+                @websocket("/ws")
+                def fake_ws(self) -> None:
+                    ...
+
             FakeController.create_router()
             apirouter.assert_called_once_with(prefix="/test", dependencies=None, tags=None)
             apirouter.return_value.add_api_route.assert_called_once_with(
@@ -88,3 +92,7 @@ def fake_method(self) -> None:
                 deprecated=True,
                 methods=["GET"],
             )
+            apirouter.return_value.add_api_websocket_route.assert_called_once_with(
+                "/ws",
+                FakeController.fake_ws,
+            )