├── tests ├── __init__.py ├── example_app │ ├── __init__.py │ ├── feature_0 │ │ ├── __init__.py │ │ └── subscribe.py │ ├── feature_1 │ │ ├── __init__.py │ │ └── ping.py │ ├── feature_2 │ │ ├── __init__.py │ │ └── alert.py │ ├── service.py │ └── main.py └── test_example_app.py ├── .devcontainer ├── .env.example ├── post-create-command.sh ├── initialize-command.sh ├── Dockerfile └── devcontainer.json ├── .gitignore ├── fastws ├── __init__.py ├── asyncapi.py ├── broker.py ├── docs.py ├── application.py └── routing.py ├── LICENSE ├── pyproject.toml ├── .github └── workflows │ └── test.yml └── README.md /tests/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /.devcontainer/.env.example: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /tests/example_app/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /tests/example_app/feature_0/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /tests/example_app/feature_1/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /tests/example_app/feature_2/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /.devcontainer/post-create-command.sh: -------------------------------------------------------------------------------- 1 | #!/bin/sh 2 | 3 | poetry install --with dev 4 | 5 | brew install act -------------------------------------------------------------------------------- /tests/example_app/feature_1/ping.py: -------------------------------------------------------------------------------- 1 | from fastws import OperationRouter 2 | 3 | router = OperationRouter(prefix="feature_1.") 4 | 5 | 6 | @router.send("ping", reply="pong") 7 | async def send_ping(): 8 | return 9 | -------------------------------------------------------------------------------- /.devcontainer/initialize-command.sh: -------------------------------------------------------------------------------- 1 | #!/bin/sh 2 | 3 | ## If .env does not exist, create it 4 | if [[ ! -e ./.devcontainer/.env ]]; then 5 | sed '/Optional/,$ s/./#&/' .devcontainer/.env.example > .devcontainer/.env 6 | fi -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | .idea 2 | .ipynb_checkpoints 3 | __pycache__ 4 | .pytest_cache 5 | htmlcov 6 | dist 7 | site 8 | .coverage 9 | coverage.xml 10 | test.db 11 | log.txt 12 | Pipfile.lock 13 | env3.* 14 | env 15 | docs_build 16 | venv 17 | .env 18 | *.log 19 | poetry.lock 20 | .vscode -------------------------------------------------------------------------------- /tests/example_app/feature_2/alert.py: -------------------------------------------------------------------------------- 1 | import logging 2 | from fastws import OperationRouter 3 | from pydantic import BaseModel 4 | 5 | router = OperationRouter(prefix="feature_2.") 6 | 7 | 8 | class AlertPayload(BaseModel): 9 | message: str 10 | 11 | 12 | @router.recv("alert") 13 | async def alert_from_server(payload: AlertPayload) -> AlertPayload: 14 | logging.info(f"{payload.message}") 15 | return payload 16 | -------------------------------------------------------------------------------- /tests/example_app/service.py: -------------------------------------------------------------------------------- 1 | from fastws import FastWS 2 | 3 | from .feature_0 import subscribe 4 | from .feature_1 import ping 5 | from .feature_2 import alert 6 | 7 | service = FastWS(title="FastWS - Broker") 8 | 9 | service.include_router(subscribe.router) 10 | service.include_router(ping.router) 11 | service.include_router(alert.router) 12 | 13 | 14 | @service.send("ping", reply="pong") 15 | def application_ping(): 16 | return 17 | -------------------------------------------------------------------------------- /fastws/__init__.py: -------------------------------------------------------------------------------- 1 | """FastWS framework. Auto-documentation WebSockets using AsyncAPI around FastAPI.""" 2 | 3 | __version__ = "0.1.7" 4 | 5 | from .application import Client as Client 6 | from .application import FastWS as FastWS 7 | from .docs import get_asyncapi as get_asyncapi 8 | from .docs import get_asyncapi_html as get_asyncapi_html 9 | from .routing import Message as Message 10 | from .routing import Operation as Operation 11 | from .routing import OperationRouter as OperationRouter 12 | -------------------------------------------------------------------------------- /.devcontainer/Dockerfile: -------------------------------------------------------------------------------- 1 | # FROM mcr.microsoft.com/devcontainers/python:0-3.11 2 | 3 | ARG VARIANT="3.11-bullseye" 4 | FROM mcr.microsoft.com/devcontainers/python:0-${VARIANT} 5 | 6 | # Install poetry for python virtual environment 7 | ENV POETRY_VIRTUALENVS_IN_PROJECT=true 8 | 9 | RUN curl -sSL https://install.python-poetry.org | POETRY_HOME=/opt/poetry python && \ 10 | cd /usr/local/bin && \ 11 | ln -s /opt/poetry/bin/poetry 12 | 13 | # Install Node.js version: none, lts/*, 16, 14, 12, 10 14 | ARG NODE_VERSION="none" 15 | RUN if [ "${NODE_VERSION}" != "none" ]; then su vscode -c "umask 0002 && . /usr/local/share/nvm/nvm.sh && nvm install ${NODE_VERSION} 2>&1"; fi 16 | -------------------------------------------------------------------------------- /tests/example_app/main.py: -------------------------------------------------------------------------------- 1 | from contextlib import asynccontextmanager 2 | from typing import Annotated 3 | 4 | from fastapi import Depends, FastAPI, Query 5 | from fastws import Client, Message 6 | 7 | from .service import service 8 | 9 | 10 | @asynccontextmanager 11 | async def lifespan(app: FastAPI): 12 | service.setup(app) 13 | yield 14 | 15 | 16 | app = FastAPI(lifespan=lifespan) 17 | 18 | 19 | @app.websocket("/") 20 | async def fastws_stream( 21 | client: Annotated[Client, Depends(service.manage)], 22 | topics: list[str] = Query([], alias="topic"), 23 | ): 24 | for t in topics: 25 | client.subscribe(t) 26 | await service.serve(client) 27 | 28 | 29 | @app.post("/{to_topic}") 30 | async def fastws_alert(to_topic: str): 31 | await service.server_send( 32 | Message(type="feature_2.alert", payload={"message": "foobar"}), 33 | topic=to_topic, 34 | ) 35 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | The MIT License (MIT) 2 | 3 | Copyright (c) 2023 Endre Krohn 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in 13 | all copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN 21 | THE SOFTWARE. -------------------------------------------------------------------------------- /tests/example_app/feature_0/subscribe.py: -------------------------------------------------------------------------------- 1 | import logging 2 | from fastws import Client, OperationRouter, FastWS 3 | from pydantic import BaseModel 4 | 5 | router = OperationRouter(prefix="feature_0.") 6 | 7 | 8 | class SubscriptionPayload(BaseModel): 9 | topic: str 10 | 11 | 12 | class SubscriptionResponse(BaseModel): 13 | detail: str 14 | topics: set[str] 15 | 16 | 17 | @router.send("subscribe", reply="subscribe.response") 18 | async def subscribe_to_topic( 19 | payload: SubscriptionPayload, 20 | client: Client, 21 | app: FastWS, 22 | ) -> SubscriptionResponse: 23 | """ 24 | Subscribe to a topic. 25 | """ 26 | client.subscribe(payload.topic) 27 | logging.info(f"app now has clients: {app.connections}") 28 | return SubscriptionResponse( 29 | detail=f"Subscribed to {payload.topic}", topics=client.topics 30 | ) 31 | 32 | 33 | @router.send("unsubscribe", reply="unsubscribe.response") 34 | async def unsubscribe_from_topic( 35 | payload: SubscriptionPayload, 36 | client: Client, 37 | ) -> SubscriptionResponse: 38 | """ 39 | Unsubscribe from a topic. 40 | """ 41 | client.unsubscribe(payload.topic) 42 | return SubscriptionResponse( 43 | detail=f"Unubscribed to {payload.topic}", topics=client.topics 44 | ) 45 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [tool.poetry] 2 | name = "fastws" 3 | version = "0.1.7" 4 | description = "FastWS framework. A WebSocket wrapper around FastAPI with auto-documentation using AsyncAPI." 5 | authors = ["Endre Krohn "] 6 | readme = "README.md" 7 | packages = [{ include = "fastws" }] 8 | repository = "https://github.com/endrekrohn/fastws" 9 | documentation = "https://github.com/endrekrohn/fastws" 10 | keywords = ["fastapi", "pydantic", "starlette", "websockets", "asyncapi"] 11 | classifiers = [ 12 | "Development Status :: 3 - Alpha", 13 | "Programming Language :: Python", 14 | "Programming Language :: Python :: 3.11", 15 | "Programming Language :: Python :: 3 :: Only", 16 | "License :: OSI Approved :: MIT License", 17 | "Operating System :: POSIX", 18 | "Intended Audience :: Developers", 19 | "Topic :: Software Development", 20 | "Topic :: Software Development :: Libraries", 21 | "Framework :: AsyncIO", 22 | "Framework :: FastAPI", 23 | "Framework :: Pydantic", 24 | ] 25 | 26 | [tool.poetry.dependencies] 27 | python = "^3.8" 28 | fastapi = ">=0.100.0" 29 | 30 | 31 | [tool.poetry.group.dev.dependencies] 32 | ipykernel = "^6.24.0" 33 | black = "^23.3.0" 34 | ruff = "^0.1.9" 35 | uvicorn = { extras = ["standard"], version = "^0.22.0" } 36 | pytest = "^7.4.0" 37 | httpx = "^0.24.1" 38 | 39 | [build-system] 40 | requires = ["poetry-core"] 41 | build-backend = "poetry.core.masonry.api" 42 | -------------------------------------------------------------------------------- /fastws/asyncapi.py: -------------------------------------------------------------------------------- 1 | from enum import Enum 2 | 3 | from pydantic import BaseModel, Field 4 | 5 | 6 | class Tag(BaseModel): 7 | name: str | Enum 8 | 9 | 10 | class Message(BaseModel): 11 | messageId: str 12 | name: str 13 | title: str 14 | summary: str | None = None 15 | description: str 16 | contentType: str = "application/json" 17 | payload: dict | None = None 18 | x_response: dict | None = Field(default=None, alias="x-response") 19 | tags: list[Tag] | None = None 20 | examples: list[dict] | None = None 21 | 22 | 23 | class Operation(BaseModel): 24 | operationId: str 25 | summary: str 26 | message: dict 27 | 28 | 29 | class Channel(BaseModel): 30 | subscribe: Operation | None = None 31 | publish: Operation | None = None 32 | 33 | 34 | class Components(BaseModel): 35 | schemas: dict 36 | messages: dict[str, Message] 37 | 38 | 39 | class ServerVariable(BaseModel): 40 | description: str | None = None 41 | default: str 42 | enum: list[str] | None = None 43 | examples: list[str] | None = None 44 | 45 | 46 | class Server(BaseModel): 47 | url: str 48 | description: str 49 | protocol: str 50 | protocolVersion: str | None = None 51 | variables: dict[str, ServerVariable] | None = None 52 | 53 | 54 | class Contact(BaseModel): 55 | name: str 56 | url: str 57 | email: str 58 | 59 | 60 | class License(BaseModel): 61 | name: str 62 | url: str | None = None 63 | 64 | 65 | class Info(BaseModel): 66 | title: str 67 | version: str 68 | description: str 69 | termsOfService: str | None = None 70 | contact: Contact | None = None 71 | license: License | None = None 72 | 73 | 74 | class AsyncAPI(BaseModel): 75 | asyncapi: str = "2.4.0" 76 | info: Info 77 | servers: dict[str, Server] | None = None 78 | channels: dict[str, Channel] 79 | components: Components 80 | defaultContentType: str = "application/json" 81 | externalDocs: dict | None = None 82 | -------------------------------------------------------------------------------- /tests/test_example_app.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | from fastapi import WebSocketDisconnect 3 | from fastapi.testclient import TestClient 4 | 5 | from tests.example_app.main import app 6 | 7 | 8 | @pytest.fixture(scope="function") 9 | def client(): 10 | with TestClient(app) as client: 11 | yield client 12 | 13 | 14 | def test_read_asyncapi_docs(client: TestClient): 15 | response = client.get("/asyncapi.json") 16 | assert response.status_code == 200 17 | response = client.get("/asyncapi") 18 | assert response.status_code == 200 19 | 20 | 21 | def test_websocket(client: TestClient): 22 | with client.websocket_connect("/?topic=foo&topic=bar") as websocket: 23 | websocket.send_json({"type": "ping"}, mode="text") 24 | data = websocket.receive_json() 25 | assert data == {"type": "pong"} 26 | 27 | websocket.send_json( 28 | {"type": "feature_0.subscribe", "payload": {"topic": "foobar"}}, 29 | mode="text", 30 | ) 31 | data = websocket.receive_json() 32 | assert "type" in data and data["type"] == "feature_0.subscribe.response" 33 | assert "payload" in data and "topics" in data["payload"] 34 | assert set(data["payload"]["topics"]) == set(["foo", "foobar", "bar"]) 35 | 36 | websocket.send_json({"type": "feature_1.ping"}, mode="text") 37 | data = websocket.receive_json() 38 | assert data == {"type": "feature_1.pong"} 39 | 40 | response = client.post("/foobar") 41 | assert response.status_code == 200 42 | 43 | data = websocket.receive_json() 44 | assert data == {"type": "feature_2.alert", "payload": {"message": "foobar"}} 45 | 46 | websocket.send_json( 47 | {"type": "feature_0.not_an_operation", "payload": {"foo": "bar"}}, 48 | mode="text", 49 | ) 50 | with pytest.raises(WebSocketDisconnect, match="No matching type") as exc_info: 51 | websocket.receive_json() 52 | assert exc_info.value.code == 1003 53 | -------------------------------------------------------------------------------- /.github/workflows/test.yml: -------------------------------------------------------------------------------- 1 | name: Test 2 | 3 | on: 4 | push: 5 | branches: 6 | - main 7 | pull_request: 8 | types: 9 | - opened 10 | - synchronize 11 | 12 | jobs: 13 | lint: 14 | runs-on: ubuntu-latest 15 | steps: 16 | - uses: actions/checkout@v3 17 | - uses: actions/setup-python@v4 18 | with: 19 | python-version: "3.11" 20 | - name: Install Poetry 21 | uses: abatilo/actions-poetry@v2 22 | - name: Setup a local venv 23 | run: | 24 | poetry config virtualenvs.create true --local 25 | poetry config virtualenvs.in-project true --local 26 | - uses: actions/cache@v3 27 | id: cache 28 | name: Define a cache for the venv based on the lock file 29 | with: 30 | path: ./.venv 31 | key: venv-${{ hashFiles('poetry.lock') }} 32 | - name: Install the project dependencies 33 | if: steps.cache.outputs.cache-hit != 'true' 34 | run: poetry install --no-cache 35 | - name: Run lint 36 | run: | 37 | poetry run ruff check . 38 | poetry run black --check . 39 | 40 | test: 41 | runs-on: ubuntu-latest 42 | steps: 43 | - uses: actions/checkout@v3 44 | - uses: actions/setup-python@v4 45 | with: 46 | python-version: "3.11" 47 | - name: Install Poetry 48 | uses: abatilo/actions-poetry@v2 49 | - name: Setup a local venv 50 | run: | 51 | poetry config virtualenvs.create true --local 52 | poetry config virtualenvs.in-project true --local 53 | - uses: actions/cache@v3 54 | id: cache 55 | name: Define a cache for the venv based on the lock file 56 | with: 57 | path: ./.venv 58 | key: venv-${{ hashFiles('poetry.lock') }} 59 | - name: Install the project dependencies 60 | if: steps.cache.outputs.cache-hit != 'true' 61 | run: poetry install --no-cache 62 | - name: Run tests 63 | run: poetry run pytest -v -------------------------------------------------------------------------------- /.devcontainer/devcontainer.json: -------------------------------------------------------------------------------- 1 | // For format details, see https://aka.ms/devcontainer.json. For config options, see the 2 | // README at: https://github.com/devcontainers/templates/tree/main/src/python 3 | { 4 | "name": "Python 3", 5 | // Or use a Dockerfile or Docker Compose file. More info: https://containers.dev/guide/dockerfile 6 | // "image": "mcr.microsoft.com/devcontainers/python:0-3.11", 7 | "build": { 8 | "dockerfile": "Dockerfile", 9 | "args": { 10 | "VARIANT": "3.11-bullseye", 11 | "NODE_VERSION": "lts/*" 12 | } 13 | }, 14 | "features": { 15 | "ghcr.io/devcontainers/features/docker-in-docker:2": { 16 | "moby": true, 17 | "azureDnsAutoDetection": true, 18 | "installDockerBuildx": true, 19 | "version": "latest", 20 | "dockerDashComposeVersion": "v2" 21 | }, 22 | "ghcr.io/meaningful-ooo/devcontainer-features/homebrew:2": {} 23 | }, 24 | // Features to add to the dev container. More info: https://containers.dev/features. 25 | // "features": {}, 26 | // Use 'forwardPorts' to make a list of ports inside the container available locally. 27 | // "forwardPorts": [], 28 | // Use 'postCreateCommand' to run commands after the container is created. 29 | "postCreateCommand": "bash .devcontainer/post-create-command.sh", 30 | "initializeCommand": "bash .devcontainer/initialize-command.sh", 31 | "runArgs": [ 32 | "--env-file", 33 | ".devcontainer/.env" 34 | ], 35 | // Configure tool-specific properties. 36 | "customizations": { 37 | "vscode": { 38 | "settings": { 39 | "window.title": "${rootName}${separator}${dirty}${activeEditorShort}${separator}${profileName}${separator}${appName}", 40 | "window.commandCenter": true, 41 | "window.menuBarVisibility": "toggle", 42 | "[json]": { 43 | "editor.defaultFormatter": "vscode.json-language-features", 44 | "editor.formatOnSave": true 45 | }, 46 | "[jsonc]": { 47 | "editor.quickSuggestions": { 48 | "strings": true 49 | }, 50 | "editor.suggest.insertMode": "replace", 51 | "editor.defaultFormatter": "vscode.json-language-features", 52 | "editor.formatOnSave": true 53 | }, 54 | "[markdown]": { 55 | "editor.defaultFormatter": "yzhang.markdown-all-in-one", 56 | "editor.formatOnSave": true 57 | }, 58 | "[python]": { 59 | "editor.defaultFormatter": "ms-python.black-formatter", 60 | "editor.formatOnSave": true, 61 | "editor.codeActionsOnSave": { 62 | "source.organizeImports": false, 63 | "source.": true 64 | } 65 | }, 66 | "[html]": { 67 | "editor.suggest.insertMode": "replace", 68 | "editor.formatOnSave": true, 69 | "editor.defaultFormatter": "vscode.html-language-features" 70 | }, 71 | "[toml]": { 72 | "editor.formatOnSave": true, 73 | "editor.defaultFormatter": "tamasfe.even-better-toml" 74 | }, 75 | "python.poetryPath": "/usr/local/bin/poetry", 76 | "python.pythonPath": "${workspaceFolder}/.venv/bin/python", 77 | "python.defaultInterpreterPath": "${workspaceFolder}/.venv/bin/python", 78 | "python.analysis.typeCheckingMode": "basic", 79 | "black-formatter.interpreter": [ 80 | "${workspaceFolder}/.venv/bin/python" 81 | ], 82 | "python.testing.pytestEnabled": true, 83 | "python.terminal.activateEnvironment": false, 84 | "jupyter.notebookFileRoot": "${workspaceFolder}/", 85 | "files.exclude": { 86 | "**/__pycache__": true, 87 | "**/*.egg-info": true, 88 | "**/.pytest_cache": true, 89 | "**/.benchmarks": true, 90 | "**/.ruff_cache": true, 91 | "**/.venv": true, 92 | "pyrepo": true 93 | }, 94 | "git.branchProtection": [ 95 | "main", 96 | "main/*" 97 | ] 98 | }, 99 | "extensions": [ 100 | "ms-python.vscode-pylance", 101 | "ms-python.python", 102 | "ms-python.black-formatter", 103 | "ms-azuretools.vscode-docker", 104 | "ms-toolsai.jupyter", 105 | "charliermarsh.ruff", 106 | "eamodio.gitlens", 107 | "tamasfe.even-better-toml", 108 | "yzhang.markdown-all-in-one" 109 | ], 110 | "recommendations": [ 111 | "GitHub.copilot" 112 | ] 113 | } 114 | } 115 | // Uncomment to connect as root instead. More info: https://aka.ms/dev-containers-non-root. 116 | // "remoteUser": "root" 117 | } -------------------------------------------------------------------------------- /fastws/broker.py: -------------------------------------------------------------------------------- 1 | import asyncio 2 | import typing 3 | from enum import Enum 4 | 5 | from fastws.docs import get_asyncapi 6 | from fastws.routing import ( 7 | Message, 8 | MessageT, 9 | MethodT, 10 | NoMatchingOperation, 11 | Operation, 12 | OperationRouter, 13 | ) 14 | 15 | 16 | async def run_handler_function( 17 | *, 18 | handler: typing.Callable[..., typing.Any], 19 | values: dict[str, typing.Any], 20 | ) -> typing.Any: 21 | if asyncio.iscoroutinefunction(handler): 22 | return await handler(**values) 23 | else: 24 | return handler(**values) 25 | 26 | 27 | class Broker: 28 | def __init__( 29 | self, 30 | title: str = "Event Driven Broker", 31 | version: str = "0.0.1", 32 | asyncapi_version: str = "2.4.0", 33 | description: str | None = None, 34 | terms_of_service: str | None = None, 35 | contact: dict[str, str] | None = None, 36 | license_info: dict[str, str] | None = None, 37 | servers: dict | None = None, 38 | ) -> None: 39 | self.title = title 40 | self.version = version 41 | self.asyncapi_version = asyncapi_version 42 | self.description = description 43 | self.terms_of_service = terms_of_service 44 | self.contact = contact 45 | self.license_info = license_info 46 | self.servers = servers 47 | self.router = OperationRouter() 48 | self.asyncapi_schema: dict | None = None 49 | 50 | def include_router( 51 | self, 52 | router: OperationRouter, 53 | *, 54 | prefix: str = "", 55 | ): 56 | self.router.include_router(router, prefix=prefix) 57 | 58 | def send( 59 | self, 60 | operation: str, 61 | name: str | None = None, 62 | tags: list[str | Enum] | None = None, 63 | summary: str | None = None, 64 | description: str | None = None, 65 | reply: str | None = None, 66 | ) -> typing.Callable[ 67 | [typing.Callable[..., typing.Any]], typing.Callable[..., typing.Any] 68 | ]: 69 | def decorator( 70 | func: typing.Callable[..., typing.Any], 71 | ) -> typing.Callable[..., typing.Any]: 72 | self.router.add_route( 73 | operation=operation, 74 | handler=func, 75 | method="SEND", 76 | name=name, 77 | tags=tags, 78 | summary=summary, 79 | description=description, 80 | reply_operation=reply, 81 | ) 82 | return func 83 | 84 | return decorator 85 | 86 | def recv( 87 | self, 88 | operation: str, 89 | name: str | None = None, 90 | tags: list[str | Enum] | None = None, 91 | summary: str | None = None, 92 | description: str | None = None, 93 | ) -> typing.Callable[ 94 | [typing.Callable[..., typing.Any]], typing.Callable[..., typing.Any] 95 | ]: 96 | def decorator( 97 | func: typing.Callable[..., typing.Any], 98 | ) -> typing.Callable[..., typing.Any]: 99 | self.router.add_route( 100 | operation=operation, 101 | handler=func, 102 | method="RECEIVE", 103 | name=name, 104 | tags=tags, 105 | summary=summary, 106 | description=description, 107 | ) 108 | return func 109 | 110 | return decorator 111 | 112 | def _match_route(self, operation: str, method: MethodT) -> Operation: 113 | for route in self.router.routes: 114 | if route.matches(operation, method): 115 | return route 116 | raise NoMatchingOperation("no matching route found") 117 | 118 | async def __call__( 119 | self, 120 | message: Message, 121 | method: MethodT = "SEND", 122 | **params, 123 | ) -> MessageT | None: 124 | route = self._match_route(operation=message.type, method=method) 125 | values = route.convert_params(message=message, params=params) 126 | result = await run_handler_function( 127 | handler=route.handler, 128 | values=values, 129 | ) 130 | if route.reply_operation is None: 131 | return 132 | return route.reply_payload.model_validate( 133 | {"type": route.reply_operation or route.operation, "payload": result} 134 | ) 135 | 136 | def asyncapi(self) -> dict[str, typing.Any]: 137 | if not self.asyncapi_schema: 138 | self.asyncapi_schema = get_asyncapi( 139 | operations=self.router.routes, 140 | title=self.title, 141 | version=self.version, 142 | asyncapi_version=self.asyncapi_version, 143 | description=self.description, 144 | terms_of_service=self.terms_of_service, 145 | contact=self.contact, 146 | license_info=self.license_info, 147 | servers=self.servers, 148 | ) 149 | return self.asyncapi_schema 150 | -------------------------------------------------------------------------------- /fastws/docs.py: -------------------------------------------------------------------------------- 1 | from dataclasses import dataclass 2 | from typing import Any, Hashable, Literal, Sequence 3 | 4 | from pydantic import BaseModel 5 | from pydantic.json_schema import GenerateJsonSchema 6 | from pydantic_core import CoreSchema 7 | 8 | from fastws import asyncapi 9 | from fastws.routing import Operation 10 | 11 | REF_SCHEMAS_TEMPLATE = "#/components/schemas/{model}" 12 | REF_MESSAGES_TEMPLATE = "#/components/messages/{message}" 13 | 14 | 15 | @dataclass 16 | class Field: 17 | key: Hashable 18 | json_mode: Literal["validation", "serialization"] 19 | core_schema: CoreSchema 20 | 21 | 22 | def get_fields( 23 | routes: Sequence[Operation], 24 | ) -> list[Field]: 25 | fields: list[Field] = [] 26 | for route in routes: 27 | if route.payload is not None and issubclass(route.payload, BaseModel): 28 | fields.append( 29 | Field( 30 | key=route.operation, 31 | json_mode="validation", 32 | core_schema=route.payload.__pydantic_core_schema__, 33 | ) 34 | ) 35 | if route.reply_payload is not None and issubclass( 36 | route.reply_payload, BaseModel 37 | ): 38 | key = route.operation 39 | if route.method == "SEND": 40 | key = route.reply_operation 41 | fields.append( 42 | Field( 43 | key=key, 44 | json_mode="validation", 45 | core_schema=route.reply_payload.__pydantic_core_schema__, 46 | ) 47 | ) 48 | return fields 49 | 50 | 51 | def get_messages( 52 | routes: Sequence[Operation], 53 | field_mapping: dict[tuple[Hashable, Literal["validation", "serialization"]], dict], 54 | ) -> tuple[dict[str, asyncapi.Message], list[str], list[str]]: 55 | messages = {} 56 | sub_messages = [] 57 | pub_messages = [] 58 | 59 | for route in routes: 60 | msg = asyncapi.Message( 61 | messageId=route.operation, 62 | name=route.name, 63 | title=" ".join(route.name.split("_")).title(), 64 | summary=route.summary, 65 | description=route.description, 66 | tags=[asyncapi.Tag(name=t) for t in route.tags], 67 | ) 68 | if route.method == "SEND": 69 | key = route.operation 70 | pub_messages.append(key) 71 | 72 | to_update = { 73 | "messageId": key, 74 | "payload": field_mapping.get((key, "validation"), None), 75 | } 76 | if route.reply_operation is not None: 77 | to_update["x_response"] = { 78 | "$ref": REF_MESSAGES_TEMPLATE.format(message=route.reply_operation) 79 | } 80 | messages[key] = msg.model_copy(update=to_update) 81 | if route.response_model or route.reply_operation or route.method == "RECEIVE": 82 | key = route.operation 83 | if route.method == "SEND": 84 | key = route.reply_operation 85 | sub_messages.append(key) 86 | messages[key] = msg.model_copy( 87 | update={ 88 | "messageId": key, 89 | "payload": field_mapping.get((key, "validation"), None), 90 | } 91 | ) 92 | return messages, sub_messages, pub_messages 93 | 94 | 95 | def get_asyncapi( 96 | operations: Sequence[Operation], 97 | title: str = "Event Driven Broker", 98 | version: str = "1.0.0", 99 | asyncapi_version: str = "2.4.0", 100 | description: str | None = None, 101 | terms_of_service: str | None = None, 102 | contact: dict[str, str] | None = None, 103 | license_info: dict[str, str] | None = None, 104 | servers: dict | None = None, 105 | ) -> dict[str, Any]: 106 | output: dict[str, Any] = {"asyncapi": asyncapi_version} 107 | 108 | output["info"] = { 109 | "title": title, 110 | "version": version, 111 | "description": description or "", 112 | "termsOfService": terms_of_service, 113 | "contact": contact, 114 | "license": license_info, 115 | } 116 | if servers is not None: 117 | output["servers"] = servers 118 | 119 | schema_generator = GenerateJsonSchema(ref_template=REF_SCHEMAS_TEMPLATE) 120 | 121 | fields = get_fields(operations) 122 | field_mapping, definitions = schema_generator.generate_definitions( 123 | inputs=[(f.key, f.json_mode, f.core_schema) for f in fields] 124 | ) 125 | messages, sub_messages, pub_messages = get_messages(operations, field_mapping) 126 | 127 | output["channels"] = { 128 | "/": { 129 | "publish": { 130 | "operationId": "sendMessage", 131 | "summary": "The API user can send a given message to the server.", 132 | "message": { 133 | "oneOf": [ 134 | {"$ref": REF_MESSAGES_TEMPLATE.format(message=v)} 135 | for v in pub_messages 136 | ] 137 | }, 138 | }, 139 | "subscribe": { 140 | "operationId": "processMessage", 141 | "summary": "The API user can receive a given message from the server.", 142 | "message": { 143 | "oneOf": [ 144 | {"$ref": REF_MESSAGES_TEMPLATE.format(message=v)} 145 | for v in sub_messages 146 | ] 147 | }, 148 | }, 149 | } 150 | } 151 | messages = { 152 | k: v.model_dump(by_alias=True, exclude_unset=True) for k, v in messages.items() 153 | } 154 | output["components"] = {"schemas": definitions, "messages": messages} 155 | return asyncapi.AsyncAPI(**output).model_dump(by_alias=True, exclude_none=True) 156 | 157 | 158 | def get_asyncapi_html( 159 | *, 160 | title: str = "AsyncAPI", 161 | asyncapi_url: str = "/asyncapi.json", 162 | asyncapi_js_url: str = "https://unpkg.com/@asyncapi/react-component@1.0.0-next.39/browser/standalone/index.js", 163 | asyncapi_css_url: str = "https://unpkg.com/@asyncapi/react-component@1.0.0-next.39/styles/default.min.css", 164 | ) -> str: 165 | html = f""" 166 | 167 | 168 | 169 | 170 | 177 | {title} 178 | 179 | 180 |
181 | 182 | 195 | 196 | """ 197 | return html 198 | -------------------------------------------------------------------------------- /fastws/application.py: -------------------------------------------------------------------------------- 1 | import asyncio 2 | import logging 3 | from typing import AsyncGenerator, AsyncIterator, Callable, Awaitable 4 | from uuid import uuid4 5 | from fastapi import Request, WebSocketException, status, FastAPI 6 | from fastapi.responses import HTMLResponse, JSONResponse 7 | from pydantic import BaseModel, ValidationError 8 | from starlette.websockets import WebSocket 9 | 10 | from fastws.routing import Message, NoMatchingOperation 11 | from fastws.broker import Broker 12 | from fastws.docs import get_asyncapi_html 13 | 14 | 15 | class Client: 16 | def __init__(self, ws: WebSocket) -> None: 17 | self.ws = ws 18 | self.uid = uuid4().hex 19 | self.topics: set[str] = set() 20 | 21 | async def send(self, message: str) -> None: 22 | await self.ws.send_text(message) 23 | 24 | def subscribe(self, topic: str) -> None: 25 | if topic not in self.topics: 26 | self.topics.add(topic) 27 | 28 | def unsubscribe(self, topic: str) -> None: 29 | if topic in self.topics: 30 | self.topics.remove(topic) 31 | 32 | async def __aiter__(self) -> AsyncIterator[Message]: 33 | async for message in self.ws.iter_text(): 34 | yield Message.model_validate_json(message) 35 | 36 | 37 | class FastWS(Broker): 38 | def __init__( 39 | self, 40 | *, 41 | title: str = "FastWS", 42 | version: str = "0.0.1", 43 | asyncapi_version: str = "2.4.0", 44 | description: str | None = None, 45 | terms_of_service: str | None = None, 46 | contact: dict[str, str] | None = None, 47 | license_info: dict[str, str] | None = None, 48 | servers: dict | None = None, 49 | asyncapi_url: str | None = "/asyncapi.json", 50 | asyncapi_docs_url: str | None = "/asyncapi", 51 | debug: bool = False, 52 | heartbeat_interval: float | None = None, 53 | max_connection_lifespan: float | None = None, 54 | auth_handler: Callable[[WebSocket], Awaitable[bool]] | None = None, 55 | auto_ws_accept: bool = True, 56 | ) -> None: 57 | super().__init__( 58 | title=title, 59 | version=version, 60 | asyncapi_version=asyncapi_version, 61 | description=description, 62 | terms_of_service=terms_of_service, 63 | contact=contact, 64 | license_info=license_info, 65 | servers=servers, 66 | ) 67 | self.connections: dict[str, Client] = {} 68 | self.debug = debug 69 | self.heartbeat_interval = heartbeat_interval 70 | self.shutdown_event = asyncio.Event() 71 | self.max_connection_lifespan = max_connection_lifespan 72 | self.auth_handler = auth_handler 73 | self.auto_ws_accept = auto_ws_accept 74 | self.asyncapi_url = asyncapi_url 75 | self.asyncapi_docs_url = asyncapi_docs_url 76 | 77 | def log(self, msg: str) -> None: 78 | if self.debug: 79 | logging.debug(f"{msg} ({len(self.connections)} conns)") 80 | 81 | def _connect(self, client: Client) -> None: 82 | self.connections[client.uid] = client 83 | self.log(f"{client} connected") 84 | 85 | def _disconnect(self, client: Client | str) -> Client | None: 86 | if isinstance(client, Client): 87 | client = client.uid 88 | return self.connections.pop(client, None) 89 | 90 | async def _auth(self, ws: WebSocket) -> bool: 91 | if self.auth_handler is None: 92 | if self.auto_ws_accept: 93 | await ws.accept() 94 | return True 95 | return await self.auth_handler(ws) 96 | 97 | async def manage(self, ws: WebSocket) -> AsyncGenerator[Client, None]: 98 | if not await self._auth(ws): 99 | return 100 | client = Client(ws) 101 | self._connect(client) 102 | try: 103 | yield client 104 | except Exception: 105 | self.log("unknown disconnect") 106 | finally: 107 | self._disconnect(client) 108 | 109 | async def broadcast(self, topic: str, message: BaseModel): 110 | msg = message.model_dump_json() 111 | async with asyncio.TaskGroup() as tg: 112 | for client in filter( 113 | lambda x: topic in x.topics, 114 | self.connections.values(), 115 | ): 116 | tg.create_task(client.send(msg)) 117 | self.log(f"Sent to {client.uid}") 118 | 119 | async def handle_exception( 120 | self, 121 | exc: ValueError | ValidationError | NoMatchingOperation | TimeoutError, 122 | ): 123 | self.log(f"could not parse message {exc}") 124 | error = WebSocketException( 125 | code=status.WS_1003_UNSUPPORTED_DATA, reason="Could not decode message" 126 | ) 127 | if isinstance(exc, ValidationError): 128 | error.reason = "Could not validate payload" 129 | if isinstance(exc, NoMatchingOperation): 130 | error.reason = "No matching type" 131 | if isinstance(exc, TimeoutError): 132 | error.reason = ( 133 | "Connection timed out. " 134 | + f"Heartbeat interval {self.heartbeat_interval or 'unset'}. " 135 | + f"Max connection lifespan {self.max_connection_lifespan or 'unset'}" 136 | ) 137 | raise error 138 | 139 | async def serve(self, client: Client): 140 | try: 141 | async with asyncio.timeout( 142 | self.heartbeat_interval 143 | ) as heartbeat_cm, asyncio.timeout( 144 | self.max_connection_lifespan, 145 | ): 146 | async for message in client: 147 | if self.heartbeat_interval is not None: 148 | heartbeat_cm.reschedule( 149 | asyncio.get_running_loop().time() + self.heartbeat_interval 150 | ) 151 | await self.client_send(message, client=client) 152 | except (ValueError, ValidationError, NoMatchingOperation, TimeoutError) as exc: 153 | await self.handle_exception(exc) 154 | 155 | async def client_send(self, message: Message, *, client: Client): 156 | if ( 157 | reply := await self(message, method="SEND", client=client, app=self) 158 | ) is not None: 159 | await client.send(reply.model_dump_json()) 160 | 161 | async def server_send(self, message: Message, *, topic: str, **params): 162 | if (reply := await self(message, method="RECEIVE", app=self, **params)) is None: 163 | return 164 | await self.broadcast(topic, reply) 165 | 166 | def setup(self, app: FastAPI) -> None: 167 | if self.asyncapi_url and self.asyncapi_docs_url: 168 | 169 | async def asyncapi_ui_html(req: Request) -> HTMLResponse: 170 | root_path = req.scope.get("root_path", "").rstrip("/") 171 | asyncapi_url = f"{root_path}{self.asyncapi_url}" 172 | return HTMLResponse( 173 | get_asyncapi_html( 174 | title=f"{self.title} - AsyncAPI UI", 175 | asyncapi_url=asyncapi_url, 176 | ) 177 | ) 178 | 179 | app.add_route( 180 | self.asyncapi_docs_url, asyncapi_ui_html, include_in_schema=False 181 | ) 182 | 183 | async def asyncapi_json(_: Request) -> JSONResponse: 184 | return JSONResponse(self.asyncapi()) 185 | 186 | app.add_route(self.asyncapi_url, asyncapi_json, include_in_schema=False) 187 | -------------------------------------------------------------------------------- /fastws/routing.py: -------------------------------------------------------------------------------- 1 | import inspect 2 | import typing 3 | from enum import Enum 4 | from typing import Generic, Literal, TypeVar, Union 5 | 6 | from fastapi.dependencies.utils import get_typed_return_annotation, get_typed_signature 7 | from pydantic import BaseModel 8 | 9 | MethodT = Literal["SEND", "RECEIVE"] 10 | PayloadT = TypeVar("PayloadT") 11 | EventTypeT = TypeVar("EventTypeT", str, int) 12 | 13 | 14 | class _Msg(BaseModel, Generic[EventTypeT]): 15 | type: EventTypeT | str 16 | 17 | 18 | class _MsgWithPayload(_Msg, Generic[EventTypeT, PayloadT]): 19 | payload: PayloadT | None 20 | 21 | 22 | class Message(_MsgWithPayload, Generic[EventTypeT, PayloadT]): 23 | payload: PayloadT | None = None 24 | 25 | 26 | MessageT = Union[_Msg, _MsgWithPayload] 27 | 28 | 29 | class NoMatchingOperation(Exception): 30 | ... 31 | 32 | 33 | def get_name(endpoint: typing.Callable) -> str: 34 | if inspect.isroutine(endpoint) or inspect.isclass(endpoint): 35 | return endpoint.__name__ 36 | return endpoint.__class__.__name__ 37 | 38 | 39 | class Operation: 40 | def __init__( 41 | self, 42 | operation: str, 43 | handler: typing.Callable[..., typing.Any], 44 | method: MethodT, 45 | *, 46 | name: str | None = None, 47 | tags: list[str | Enum] | None = None, 48 | summary: str | None = None, 49 | description: str | None = None, 50 | reply_operation: str | None = None, 51 | ) -> None: 52 | self.operation = operation 53 | self.handler = handler 54 | self.method: MethodT = method 55 | self.name = get_name(handler) if name is None else name 56 | self.tags = tags or [] 57 | self.summary = summary 58 | self.description = description or inspect.cleandoc(self.handler.__doc__ or "") 59 | self.description = self.description.split("\f")[0].strip() 60 | self.response_model = get_typed_return_annotation(self.handler) 61 | self.parameters = get_typed_signature(self.handler).parameters.copy() 62 | self.reply_operation = ( 63 | reply_operation if self.method == "SEND" else self.operation 64 | ) 65 | if self.method == "SEND" and self.response_model is not None: 66 | assert ( 67 | self.reply_operation is not None 68 | ), "Send operations with a response model defined must include a reply" 69 | 70 | op_t = Literal[self.operation] # type: ignore 71 | reply_t = Literal[self.reply_operation or self.operation] # type: ignore 72 | 73 | self.payload = _Msg[op_t] 74 | self.reply_payload = _Msg[reply_t] 75 | 76 | if (p := self.parameters.get("payload", None)) is not None: 77 | self.payload = _MsgWithPayload[op_t, p.annotation] 78 | 79 | if self.response_model is not None: 80 | self.reply_payload = _MsgWithPayload[reply_t, self.response_model] 81 | 82 | def matches(self, operation: str, method: MethodT) -> bool: 83 | return self.operation == operation and self.method == method 84 | 85 | def convert_params( 86 | self, 87 | message: Message, 88 | params: dict[str, typing.Any], 89 | ) -> dict[str, typing.Any]: 90 | converted_params = {} 91 | 92 | for k, v in self.parameters.items(): 93 | if k == "payload" and issubclass(v.annotation, BaseModel): 94 | converted_params[k] = v.annotation.model_validate(message.payload) 95 | else: 96 | if k not in params: 97 | raise RuntimeError(f"Missing parameter {k} of type {v.annotation}") 98 | converted_params[k] = params.get(k) 99 | if not (self.parameters.keys() == converted_params.keys()): 100 | raise RuntimeError("Missing parameters in function call") 101 | return converted_params 102 | 103 | 104 | class OperationRouter: 105 | def __init__( 106 | self, 107 | *, 108 | prefix: str = "", 109 | tags: list[str | Enum] | None = None, 110 | routes: list[Operation] | None = None, 111 | ) -> None: 112 | self.prefix = prefix 113 | self.tags = tags or [] 114 | self.routes = routes or [] 115 | 116 | def add_route( 117 | self, 118 | operation: str, 119 | handler: typing.Callable[..., typing.Any], 120 | method: MethodT, 121 | *, 122 | name: str | None = None, 123 | tags: list[str | Enum] | None = None, 124 | summary: str | None = None, 125 | description: str | None = None, 126 | reply_operation: str | None = None, 127 | ) -> None: 128 | existing_operations = [h.operation for h in self.routes] + [ 129 | h.reply_operation for h in self.routes if h.reply_operation is not None 130 | ] 131 | assert ( 132 | operation not in existing_operations 133 | ), f"handler with operation '{operation}' already added" 134 | assert ( 135 | reply_operation not in existing_operations 136 | ), f"handler with operation '{reply_operation}' already added" 137 | current_tags = self.tags.copy() 138 | if tags: 139 | current_tags.extend(tags) 140 | route = Operation( 141 | operation=operation, 142 | handler=handler, 143 | method=method, 144 | name=name, 145 | summary=summary, 146 | description=description, 147 | tags=current_tags, 148 | reply_operation=reply_operation, 149 | ) 150 | self.routes.append(route) 151 | 152 | def include_router( 153 | self, 154 | router: "OperationRouter", 155 | *, 156 | prefix: str = "", 157 | ): 158 | for route in router.routes: 159 | self.add_route( 160 | operation=f"{prefix}{router.prefix}{route.operation}", 161 | handler=route.handler, 162 | method=route.method, 163 | name=route.name, 164 | tags=route.tags, 165 | summary=route.summary, 166 | description=route.description, 167 | reply_operation=f"{prefix}{router.prefix}{route.reply_operation}", 168 | ) 169 | 170 | def send( 171 | self, 172 | operation: str, 173 | name: str | None = None, 174 | tags: list[str | Enum] | None = None, 175 | summary: str | None = None, 176 | description: str | None = None, 177 | reply: str | None = None, 178 | ) -> typing.Callable[ 179 | [typing.Callable[..., typing.Any]], typing.Callable[..., typing.Any] 180 | ]: 181 | def decorator( 182 | func: typing.Callable[..., typing.Any], 183 | ) -> typing.Callable[..., typing.Any]: 184 | self.add_route( 185 | operation=operation, 186 | handler=func, 187 | method="SEND", 188 | name=name, 189 | tags=tags, 190 | summary=summary, 191 | description=description, 192 | reply_operation=reply, 193 | ) 194 | return func 195 | 196 | return decorator 197 | 198 | def recv( 199 | self, 200 | operation: str, 201 | name: str | None = None, 202 | tags: list[str | Enum] | None = None, 203 | summary: str | None = None, 204 | description: str | None = None, 205 | ) -> typing.Callable[ 206 | [typing.Callable[..., typing.Any]], typing.Callable[..., typing.Any] 207 | ]: 208 | def decorator( 209 | func: typing.Callable[..., typing.Any], 210 | ) -> typing.Callable[..., typing.Any]: 211 | self.add_route( 212 | operation=operation, 213 | handler=func, 214 | method="RECEIVE", 215 | name=name, 216 | tags=tags, 217 | summary=summary, 218 | description=description, 219 | ) 220 | return func 221 | 222 | return decorator 223 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # FastWS 2 | 3 |

4 | 5 | FastWS 6 | 7 |

8 | 9 | **Source Code**: https://github.com/endrekrohn/fastws 10 | 11 | --- 12 | 13 | FastWS is a wrapper around FastAPI to create better WebSocket applications with auto-documentation using AsyncAPI, in a similar fashion as FastAPIs existing use of OpenAPI. 14 | 15 | The current supported AsyncAPI verison is `2.4.0`. Once version `3.0.0` is released the plan is to upgrade to this standard. 16 | 17 | --- 18 | ## Example project 19 | 20 | If you are familiar with FastAPI and want to look at an example project using FastWS look here👨‍💻 21 | 22 | --- 23 | 24 | ## Requirements 25 | 26 | Python 3.11+ 27 | 28 | `FastWS` uses Pydantic v2 and FastAPI. 29 | 30 | ## Installation 31 | 32 | 33 | ```console 34 | $ pip install fastws 35 | ``` 36 | 37 | 38 | You will also need an ASGI server, for production such as Uvicorn or Hypercorn. 39 | 40 |
41 | 42 | ```console 43 | $ pip install "uvicorn[standard]" 44 | ``` 45 | 46 |
47 | 48 | ## Example 49 | 50 | ### Create it 51 | 52 | * Create a file `main.py` with: 53 | 54 | ```Python 55 | from contextlib import asynccontextmanager 56 | from typing import Annotated 57 | 58 | from fastapi import Depends, FastAPI 59 | from fastws import Client, FastWS 60 | 61 | service = FastWS() 62 | 63 | 64 | @service.send("ping", reply="ping") 65 | async def send_event_a(): 66 | return 67 | 68 | 69 | @asynccontextmanager 70 | async def lifespan(app: FastAPI): 71 | service.setup(app) 72 | yield 73 | 74 | 75 | app = FastAPI(lifespan=lifespan) 76 | 77 | 78 | @app.websocket("/") 79 | async def fastws_stream(client: Annotated[Client, Depends(service.manage)]): 80 | await service.serve(client) 81 | ``` 82 | 83 | We can look at the generated documentation at `http://localhost:/asyncapi`. 84 | 85 |

86 | 87 | AsyncAPI Docs 88 | 89 |

90 | 91 | --- 92 | 93 | ### Example breakdown 94 | 95 | First we import and initialize the service. 96 | 97 | 98 | ```Python 99 | from fastws import Client, FastWS 100 | 101 | service = FastWS() 102 | ``` 103 | 104 | #### Define event 105 | 106 | Next up we connect an operation (a WebSocket message) to the service, using the decorator `@service.send(...)`. We need to define the operation using a string similar to how we define an HTTP-endpoint using a path. 107 | 108 | The operation-identificator is in this case `"ping"`, meaning we will use this string to identify what type of message we are receiving. 109 | 110 | ```Python 111 | @service.send("ping", reply="ping") 112 | async def send_event_a(): 113 | return 114 | ``` 115 | 116 | If we want to define an `payload` for the operation we can extend the example: 117 | 118 | ```Python 119 | from pydantic import BaseModel 120 | 121 | class PingPayload(BaseModel): 122 | foo: str 123 | 124 | @service.send("ping", reply="ping") 125 | async def send_event_a(payload: PingPayload): 126 | return 127 | ``` 128 | 129 | An incoming message should now have the following format. (We will later view this in the generated AsyncAPI-documentation). 130 | 131 | ```json 132 | { 133 | "type": "ping", 134 | "payload": { 135 | "foo": "bar" 136 | } 137 | } 138 | ``` 139 | #### Connect service 140 | 141 | Next up we connect the service to our running FastAPI application. 142 | 143 | ```Python 144 | @asynccontextmanager 145 | async def lifespan(app: FastAPI): 146 | service.setup(app) 147 | yield 148 | 149 | 150 | app = FastAPI(lifespan=lifespan) 151 | 152 | 153 | @app.websocket("/") 154 | async def fastws_stream(client: Annotated[Client, Depends(service.manage)]): 155 | await service.serve(client) 156 | ``` 157 | 158 | The function `service.setup(app)` inside FastAPIs lifespan registers two endpoints 159 | - `/asyncapi.json`, to retrieve our API definition 160 | - `/asyncapi`, to view the AsyncAPI documentation UI. 161 | 162 | You can override both of these URLs when initializing the service, or set them to `None` to avoid registering the endpoints at all. 163 | 164 | ## Routing 165 | 166 | To spread out our service we can use the `OperationRouter`-class. 167 | 168 | ```Python 169 | # feature_1.py 170 | from fastws import Client, OperationRouter 171 | from pydantic import BaseModel 172 | 173 | router = OperationRouter(prefix="user.") 174 | 175 | 176 | class SubscribePayload(BaseModel): 177 | topic: str 178 | 179 | 180 | class SubscribeResponse(BaseModel): 181 | detail: str 182 | topics: set[str] 183 | 184 | 185 | @router.send("subscribe", reply="subscribe.response") 186 | async def subscribe_to_topic( 187 | payload: SubscribePayload, 188 | client: Client, 189 | ) -> SubscribeResponse: 190 | client.subscribe(payload.topic) 191 | return SubscribeResponse( 192 | detail=f"Subscribed to {payload.topic}", 193 | topics=client.topics, 194 | ) 195 | ``` 196 | 197 | We can then include the router in our main service. 198 | 199 | ```Python 200 | # main.py 201 | from fastws import Client, FastWS 202 | 203 | from feature_1 import router 204 | 205 | service = FastWS() 206 | service.include_router(router) 207 | ``` 208 | 209 | ## Operations, `send` and `recv` 210 | 211 | The service enables two types of operations. Let us define these operations clearly: 212 | 213 | - `send`: An operation where API user sends a message to the API server. 214 | 215 | **Note**: Up to AsyncAPI version `2.6.0` this refers to a `publish`-operation, but is changing to `send` in version `3.0.0`. 216 | 217 | - `recv`: An operation where API server sends a message to the API user. 218 | 219 | **Note**: Up to AsyncAPI version `2.6.0` this refers to a `subscribe`-operation, but is changing to `receive` in version `3.0.0`. 220 | 221 | 222 | ### The `send`-operation 223 | 224 | The above examples have only displayed the use of `send`-operations. 225 | 226 | When using the functions `FastWS.client_send(message, client)` or `FastWS.serve(client)`, we implicitly send some arguments. These keyword-arguments have the following keywords and types: 227 | 228 | - `client` with type `fastws.application.Client` 229 | - `app` with type `fastws.application.FastWS` 230 | - `payload`, optional with type defined in the function processing the message. 231 | 232 | A `send`-operation can therefore access the following arguments: 233 | 234 | ```Python 235 | from fastws import Client, FastWS 236 | from pydantic import BaseModel 237 | 238 | class Something(BaseModel): 239 | foo: str 240 | 241 | 242 | class Thing(BaseModel): 243 | bar: str 244 | 245 | 246 | @router.send("foo", reply="bar") 247 | async def some_function( 248 | payload: Something, 249 | client: Client, 250 | app: FastWS, 251 | ) -> Thing: 252 | print(f"{app.connections=}") 253 | print(f"{client.uid=}") 254 | 255 | return Thing(bar=client.uid) 256 | ``` 257 | 258 | ### The `recv`-operation 259 | 260 | When using the function `FastWS.server_send(message, topic)`, we implicitly send some arguments. These keyword-arguments have the keywords and types: 261 | 262 | - `app` with type `fastws.application.FastWS` 263 | - Optional `payload` with type defined in the function processing the message. 264 | 265 | A `recv`-operation can therefore access the following arguments: 266 | 267 | ```Python 268 | from fastws import FastWS 269 | from pydantic import BaseModel 270 | 271 | class AlertPayload(BaseModel): 272 | message: str 273 | 274 | 275 | @router.recv("alert") 276 | async def recv_client(payload: AlertPayload, app: FastWS) -> str: 277 | return "hey there!" 278 | ``` 279 | 280 | If we want create a message on the server side we can do the following: 281 | 282 | ```Python 283 | from fastapi import FastAPI 284 | from fastws import FastWS 285 | 286 | service = FastWS() 287 | app = FastAPI() 288 | 289 | @app.post("/") 290 | async def alert_on_topic_foobar(message: str): 291 | await service.server_send( 292 | Message(type="alert", payload={"message": message}), 293 | topic="foobar", 294 | ) 295 | return "ok" 296 | ``` 297 | 298 | In the example above all connections subscribed to the topic `foobar` will recieve a message the the payload `"hey there!"`. 299 | 300 | In this way you can on the server-side choose to publish messages from anywhere to any topic. This is especially useful if you have a persistent connection to Redis or similar that reads messages from some channel and want to propagate these to your users. 301 | 302 | ## Authentication 303 | 304 | There are to ways to tackle authentication using `FastWS`. 305 | 306 | ### By defining `auth_handler` 307 | 308 | One way is to provide a custom `auth_handler` when initializing the service. Below is an example where the API user must provide a secret message within a timeout to authenticate. 309 | 310 | ```Python 311 | import asyncio 312 | import logging 313 | from fastapi import WebSocket 314 | from fastws import FastWS 315 | 316 | 317 | def custom_auth(to_wait: float = 5): 318 | async def handler(ws: WebSocket) -> bool: 319 | await ws.accept() 320 | try: 321 | initial_msg = await asyncio.wait_for( 322 | ws.receive_text(), 323 | timeout=to_wait, 324 | ) 325 | return initial_msg == "SECRET_HUSH_HUSH" 326 | except asyncio.exceptions.TimeoutError: 327 | logging.info("Took to long to provide authentication") 328 | 329 | return False 330 | 331 | return handler 332 | 333 | 334 | service = FastWS(auth_handler=custom_auth()) 335 | ``` 336 | 337 | ### By using FastAPI dependency 338 | 339 | If you want to use your own FastAPI dependency to handle authentication before it enters the FastWS service you will have to set `auto_ws_accept` to `False`. 340 | 341 | ```Python 342 | import asyncio 343 | from typing import Annotated 344 | 345 | from fastapi import Depends, FastAPI, WebSocket, WebSocketException, status 346 | from fastws import Client, FastWS 347 | 348 | service = FastWS(auto_ws_accept=False) 349 | 350 | app = FastAPI() 351 | 352 | 353 | async def custom_dep(ws: WebSocket): 354 | await ws.accept() 355 | initial_msg = await asyncio.wait_for( 356 | ws.receive_text(), 357 | timeout=5, 358 | ) 359 | if initial_msg == "SECRET_HUSH_HUSH": 360 | return 361 | raise WebSocketException( 362 | code=status.WS_1008_POLICY_VIOLATION, 363 | reason="Not authenticated", 364 | ) 365 | 366 | 367 | @app.websocket("/") 368 | async def fastws_stream( 369 | client: Annotated[Client, Depends(service.manage)], 370 | _=Depends(custom_dep), 371 | ): 372 | await service.serve(client) 373 | ``` 374 | 375 | ## Heartbeats and connection lifespan 376 | 377 | To handle a WebSocket's lifespan at an application level, FastWS tries to help you by using `asyncio.timeout()`-context managers in its `serve(client)`-function. 378 | 379 | You can set the both: 380 | - `heartbeat_interval`: Meaning a client needs to send a message within this time. 381 | - `max_connection_lifespan`: Meaning all connections will disconnect when exceeding this time. 382 | 383 | These must set during initialization: 384 | 385 | ```Python 386 | from fastws import FastWS 387 | 388 | service = FastWS( 389 | heartbeat_interval=10, 390 | max_connection_lifespan=300, 391 | ) 392 | ``` 393 | 394 | Both `heartbeat_interval` and `max_connection_lifespan` can be set to None to disable any restrictions. Note this is the default. 395 | 396 | Please note that you can also set restrictions in your ASGI-server. These restrictions apply at a protocol/server-level and are different from the restrictions set by your application. Applicable settings for [uvicorn](https://www.uvicorn.org/#command-line-options): 397 | - `--ws-ping-interval` INTEGER 398 | - `--ws-ping-timeout` INTEGER 399 | - `--ws-max-size` INTEGER --------------------------------------------------------------------------------