ws

3503e25e · Dorian Turba · 94391816 · 3503e25e · 3503e25e · 3503e25e
Commit 3503e25e authored 5 days ago by Dorian Turba
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -7,6 +7,7 @@ requires-python = ">=3.13"
 dependencies = [
    "fastapi>=0.115.11",
    "uvicorn>=0.34.0",
+    "websockets>=15.0.1",
 ]
 [build-system]

--- a/src/fastapi_v1_to_v2/__init__.py
+++ b/src/fastapi_v1_to_v2/__init__.py
+import collections
+import contextlib
+import functools
 import fastapi
+from fastapi.responses import HTMLResponse
 app = fastapi.FastAPI()
+html = """
+<!DOCTYPE html>
+<html>
+    <head>
+        <title>Chat</title>
+    </head>
+    <body>
+        <h1>WebSocket Chat</h1>
+        <form action="" onsubmit="sendMessage(event)">
+            <input type="text" id="messageText" autocomplete="off"/>
+            <button>Send</button>
+        </form>
+        <ul id='messages'>
+        </ul>
+        <script>
+            var ws = new WebSocket("ws://localhost:8000/v2/foobar");
+            ws.onmessage = function(event) {
+                var messages = document.getElementById('messages')
+                var message = document.createElement('li')
+                var content = document.createTextNode(event.data)
+                message.appendChild(content)
+                messages.appendChild(message)
+            };
+            function sendMessage(event) {
+                var input = document.getElementById("messageText")
+                ws.send(input.value)
+                input.value = ''
+                event.preventDefault()
+            }
+        </script>
+    </body>
+</html>
+"""
 @app.get("/")
-async def root():
+async def get():
-    return {"message": "Hello World"}
+    return HTMLResponse(html)
 router = fastapi.APIRouter()
@@ -26,9 +65,105 @@ app.include_router(router, prefix="/v1")
 router_v2 = fastapi.APIRouter()
-@router_v2.get("/foobar")
+class ConnectionManager:
-async def foobar_v2():
+    def __init__(self):
-    return "foobar V2"
+        self.active_connections: list[fastapi.WebSocket] = []
+    async def connect(self, websocket: fastapi.WebSocket):
+        await websocket.accept()
+        self.active_connections.append(websocket)
+    def disconnect(self, websocket: fastapi.WebSocket):
+        self.active_connections.remove(websocket)
+    async def send_personal_message(self, message: str, websocket: fastapi.WebSocket):
+        await websocket.send_text(message)
+    async def broadcast(self, message: str):
+        for connection in self.active_connections:
+            await connection.send_text(message)
+class Disconnect(Exception):
+    pass
+class ConnectionManagerV2:
+    active_connections: dict[str, list[fastapi.WebSocket]] = collections.defaultdict(list)
+    current: fastapi.WebSocket
+    context: str
+    @staticmethod
+    async def enter_hook(self):
+        pass
+    @staticmethod
+    async def exit_hook(self):
+        pass
+    def __init__(
+        self,
+        ws: fastapi.WebSocket,
+        context: str = "",
+        enter_hook=None,
+        exit_hook=None,
+    ):
+        self.context = context
+        self.current = ws
+        if enter_hook is not None:
+            self.enter_hook = enter_hook
+        if exit_hook is not None:
+            self.exit_hook = exit_hook
+    async def __aenter__(self):
+        await self.current.accept()
+        self.active_connections[self.context].append(self.current)
+        await self.enter_hook(self)
+        return self
+    async def __aexit__(self, exc_type, exc, tb):
+        self.active_connections[self.context].remove(self.current)
+        await self.exit_hook(self)
+        return True
+    async def receive_text(self):
+        try:
+            return await self.current.receive_text()
+        except fastapi.WebSocketDisconnect:
+            raise Disconnect from None
+    async def send_personal_message(self, message: str):
+        await self.current.send_text(message)
+    async def broadcast(self, message: str):
+        for connection in self.active_connections[self.context]:
+            await connection.send_text(message)
+async def chat_exit_broadcast(self):
+    await self.broadcast(f"Client left the chat")
+async def chat_enter_broadcast(self):
+    await self.broadcast(f"Client join the chat, say hello to him")
+ChatManager = functools.partial(
+    ConnectionManagerV2,
+    enter_hook=chat_enter_broadcast,
+    exit_hook=chat_exit_broadcast,
+)
+@router_v2.websocket("/foobar")
+async def foobar_v2(websocket: fastapi.WebSocket):
+    async with ChatManager(
+        websocket,
+        "foobar",
+    ) as socket:
+        while True:
+            data = await socket.receive_text()
+            await socket.send_personal_message(f"You wrote: {data}")
+            await socket.broadcast(f"Client says: {data}")
 app.include_router(router_v2, prefix="/v2")
 app.include_router(router, prefix="/v2")
\ No newline at end of file
--- a/uv.lock
+++ b/uv.lock
@@ -65,12 +65,14 @@ source = { editable = "." }
 dependencies = [
    { name = "fastapi" },
    { name = "uvicorn" },
+    { name = "websockets" },
 ]
 [package.metadata]
 requires-dist = [
    { name = "fastapi", specifier = ">=0.115.11" },
    { name = "uvicorn", specifier = ">=0.34.0" },
+    { name = "websockets", specifier = ">=15.0.1" },
 ]
 [[package]]
@@ -172,3 +174,23 @@ sdist = { url = "https://files.pythonhosted.org/packages/4b/4d/938bd85e5bf2edeec
 wheels = [
    { url = "https://files.pythonhosted.org/packages/61/14/33a3a1352cfa71812a3a21e8c9bfb83f60b0011f5e36f2b1399d51928209/uvicorn-0.34.0-py3-none-any.whl", hash = "sha256:023dc038422502fa28a09c7a30bf2b6991512da7dcdb8fd35fe57cfc154126f4", size = 62315 },
 ]
+[[package]]
+name = "websockets"
+version = "15.0.1"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/21/e6/26d09fab466b7ca9c7737474c52be4f76a40301b08362eb2dbc19dcc16c1/websockets-15.0.1.tar.gz", hash = "sha256:82544de02076bafba038ce055ee6412d68da13ab47f0c60cab827346de828dee", size = 177016 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/cb/9f/51f0cf64471a9d2b4d0fc6c534f323b664e7095640c34562f5182e5a7195/websockets-15.0.1-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:ee443ef070bb3b6ed74514f5efaa37a252af57c90eb33b956d35c8e9c10a1931", size = 175440 },
+    { url = "https://files.pythonhosted.org/packages/8a/05/aa116ec9943c718905997412c5989f7ed671bc0188ee2ba89520e8765d7b/websockets-15.0.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:5a939de6b7b4e18ca683218320fc67ea886038265fd1ed30173f5ce3f8e85675", size = 173098 },
+    { url = "https://files.pythonhosted.org/packages/ff/0b/33cef55ff24f2d92924923c99926dcce78e7bd922d649467f0eda8368923/websockets-15.0.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:746ee8dba912cd6fc889a8147168991d50ed70447bf18bcda7039f7d2e3d9151", size = 173329 },
+    { url = "https://files.pythonhosted.org/packages/31/1d/063b25dcc01faa8fada1469bdf769de3768b7044eac9d41f734fd7b6ad6d/websockets-15.0.1-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:595b6c3969023ecf9041b2936ac3827e4623bfa3ccf007575f04c5a6aa318c22", size = 183111 },
+    { url = "https://files.pythonhosted.org/packages/93/53/9a87ee494a51bf63e4ec9241c1ccc4f7c2f45fff85d5bde2ff74fcb68b9e/websockets-15.0.1-cp313-cp313-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:3c714d2fc58b5ca3e285461a4cc0c9a66bd0e24c5da9911e30158286c9b5be7f", size = 182054 },
+    { url = "https://files.pythonhosted.org/packages/ff/b2/83a6ddf56cdcbad4e3d841fcc55d6ba7d19aeb89c50f24dd7e859ec0805f/websockets-15.0.1-cp313-cp313-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:0f3c1e2ab208db911594ae5b4f79addeb3501604a165019dd221c0bdcabe4db8", size = 182496 },
+    { url = "https://files.pythonhosted.org/packages/98/41/e7038944ed0abf34c45aa4635ba28136f06052e08fc2168520bb8b25149f/websockets-15.0.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:229cf1d3ca6c1804400b0a9790dc66528e08a6a1feec0d5040e8b9eb14422375", size = 182829 },
+    { url = "https://files.pythonhosted.org/packages/e0/17/de15b6158680c7623c6ef0db361da965ab25d813ae54fcfeae2e5b9ef910/websockets-15.0.1-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:756c56e867a90fb00177d530dca4b097dd753cde348448a1012ed6c5131f8b7d", size = 182217 },
+    { url = "https://files.pythonhosted.org/packages/33/2b/1f168cb6041853eef0362fb9554c3824367c5560cbdaad89ac40f8c2edfc/websockets-15.0.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:558d023b3df0bffe50a04e710bc87742de35060580a293c2a984299ed83bc4e4", size = 182195 },
+    { url = "https://files.pythonhosted.org/packages/86/eb/20b6cdf273913d0ad05a6a14aed4b9a85591c18a987a3d47f20fa13dcc47/websockets-15.0.1-cp313-cp313-win32.whl", hash = "sha256:ba9e56e8ceeeedb2e080147ba85ffcd5cd0711b89576b83784d8605a7df455fa", size = 176393 },
+    { url = "https://files.pythonhosted.org/packages/1b/6c/c65773d6cab416a64d191d6ee8a8b1c68a09970ea6909d16965d26bfed1e/websockets-15.0.1-cp313-cp313-win_amd64.whl", hash = "sha256:e09473f095a819042ecb2ab9465aee615bd9c2028e4ef7d933600a8401c79561", size = 176837 },
+    { url = "https://files.pythonhosted.org/packages/fa/a8/5b41e0da817d64113292ab1f8247140aac61cbf6cfd085d6a0fa77f4984f/websockets-15.0.1-py3-none-any.whl", hash = "sha256:f7a866fbc1e97b5c617ee4116daaa09b722101d4a3c170c787450ba409f9736f", size = 169743 },
+]