raise error if session is modified after response

nicegui/storage.py

@@ -1,20 +1,48 @@
 import contextvars
-from typing import Any, Dict
+from typing import Dict
 
 from fastapi import Request
 from starlette.middleware.base import BaseHTTPMiddleware
 
-request_contextvar = contextvars.ContextVar('request_var')
+request_contextvar = contextvars.ContextVar('request_var', default=None)
+
+
+class ReadOnlyDict:
+    def __init__(self, data: Dict, write_error_message: str = 'Read-only dict'):
+        self._data = data
+        self._write_error_message = write_error_message
+
+    def __getitem__(self, item):
+        return self._data[item]
+
+    def __iter__(self):
+        return iter(self._data)
+
+    def __len__(self):
+        return len(self._data)
+
+    def __setitem__(self, key, value):
+        raise TypeError(self._write_error_message)
 
 
 class RequestTrackingMiddleware(BaseHTTPMiddleware):
     async def dispatch(self, request: Request, call_next):
-        request_contextvar.set(request)
-        return await call_next(request)
+        token = request_contextvar.set(request)
+        request.state.responded = False
+        response = await call_next(request)
+        request.state.responded = True
+        request_contextvar.reset(token)
+        return response
 
 
 class Storage:
 
     @property
     def session(self) -> Dict:
-        return request_contextvar.get().session
+        request: Request = request_contextvar.get()
+        if request.state.responded:
+            return ReadOnlyDict(
+                request.session,
+                'the response to the browser has already been build so modifications can not be send back anymore'
+            )
+        return request.session

tests/test_storage.py

@@ -43,12 +43,13 @@ def test_session_storage_supports_asyncio(screen: Screen):
 
 
 def test_session_modifications_after_page_load(screen: Screen):
-    random_data = str(uuid4())
-
     @ui.page('/')
     async def page(client: Client):
         await client.connected()
-        app.storage.session['test'] = random_data
+        try:
+            app.storage.session['test'] = 'data'
+        except TypeError as e:
+            ui.label(str(e))
 
     screen.open('/')
-    screen.should_contain(random_data)
+    screen.should_contain('response to the browser has already been build')