Merge pull request #563 from callmephilip/incl-session-param-in-ws

Incl session param in ws

Author: jph00

fasthtml/core.py

@@ -239,6 +239,7 @@ def _find_wsp(ws, data, hdrs, arg:str, p:Parameter):
         if arg.lower()=='htmx': return _get_htmx(hdrs)
         if arg.lower()=='app': return ws.scope['app']
         if arg.lower()=='send': return partial(_send_ws, ws)
+        if 'session'.startswith(arg.lower()): return ws.scope.get('session', {})
         return None
     res = data.get(arg, None)
     if res is empty or res is None: res = hdrs.get(arg, None)

nbs/api/00_core.ipynb

@@ -796,6 +796,7 @@
     "        if arg.lower()=='htmx': return _get_htmx(hdrs)\n",
     "        if arg.lower()=='app': return ws.scope['app']\n",
     "        if arg.lower()=='send': return partial(_send_ws, ws)\n",
+    "        if 'session'.startswith(arg.lower()): return ws.scope.get('session', {})\n",
     "        return None\n",
     "    res = data.get(arg, None)\n",
     "    if res is empty or res is None: res = hdrs.get(arg, None)\n",
@@ -1228,7 +1229,7 @@
     {
      "data": {
       "text/plain": [
-       "'a604e4a2-08e8-462d-aff9-15468891fe09'"
+       "'77486da1-c613-48be-80c4-9cae89eeec48'"
       ]
      },
      "execution_count": null,
@@ -1748,17 +1749,24 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Message text was: Hi!, from client: Address(host='testclient', port=50000)\n"
+      "Message text was: Hi! with session bar, from client: Address(host='testclient', port=50000)\n"
      ]
     }
    ],
    "source": [
+    "@app.get(\"/setsess\")\n",
+    "def set_sess(session):\n",
+    "   session['foo'] = 'bar'\n",
+    "   return 'ok'\n",
+    "\n",
     "@app.ws(\"/ws\")\n",
-    "def ws(self, msg:str, ws:WebSocket): return f\"Message text was: {msg}, from client: {ws.client}\"\n",
+    "def ws(self, msg:str, ws:WebSocket, session): return f\"Message text was: {msg} with session {session.get('foo')}, from client: {ws.client}\"\n",
+    "\n",
+    "cli.get('/setsess')\n",
     "with cli.websocket_connect('/ws') as ws:\n",
     "    ws.send_text('{\"msg\":\"Hi!\"}')\n",
     "    data = ws.receive_text()\n",
-    "assert 'Message text was: Hi!' in data\n",
+    "assert 'Message text was: Hi! with session bar' in data\n",
     "print(data)"
    ]
   },
@@ -2415,13 +2423,13 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Set to 2024-10-28 20:22:34.772989\n"
+      "Set to 2024-11-04 15:30:23.038930\n"
      ]
     },
     {
      "data": {
       "text/plain": [
-       "'Session time: 2024-10-28 20:22:34.772989'"
+       "'Session time: 2024-11-04 15:30:23.038930'"
       ]
      },
      "execution_count": null,
@@ -2786,7 +2794,7 @@
     {
      "data": {
       "text/plain": [
-       "'Cookie was set at time 20:22:35.467691'"
+       "'Cookie was set at time 15:30:23.174646'"
       ]
      },
      "execution_count": null,

nbs/explains/websockets.ipynb

@@ -93,6 +93,33 @@
     "This is a fairly simple example and could be done just as easily with standard HTTP requests, but it illustrates the basic idea of how websockets work. Let's look at a more complex example next."
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Session data in Websockets\n",
+    "\n",
+    "Session data is shared between standard HTTP routes and Websockets. This means you can access, for example, logged in user ID inside websocket handler:\n",
+    "\n",
+    "```python\n",
+    "from fasthtml.common import *\n",
+    "\n",
+    "app = FastHTML(exts='ws')\n",
+    "rt = app.route\n",
+    "\n",
+    "@rt('/login')\n",
+    "def get(session):\n",
+    "    session[\"person\"] = \"Bob\"\n",
+    "    return \"ok\"\n",
+    "\n",
+    "@app.ws('/ws')\n",
+    "async def ws(msg:str, send, session):\n",
+    "    await send(Div(f'Hello {session.get(\"person\")}' + msg, id='notifications'))\n",
+    "\n",
+    "serve()\n",
+    "```"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},