1
0
Эх сурвалжийг харах

split user storage into separate files; fix mypy

Falko Schindler 1 жил өмнө
parent
commit
bae918679d
2 өөрчлөгдсөн 15 нэмэгдсэн , 13 устгасан
  1. 12 10
      nicegui/storage.py
  2. 3 3
      tests/test_storage.py

+ 12 - 10
nicegui/storage.py

@@ -5,7 +5,7 @@ import threading
 import uuid
 from collections.abc import MutableMapping
 from pathlib import Path
-from typing import Any, Dict, Iterator
+from typing import Any, Dict, Iterator, Optional, Union
 
 import aiofiles
 from fastapi import Request
@@ -14,7 +14,7 @@ from starlette.responses import Response
 
 from . import globals
 
-request_contextvar = contextvars.ContextVar('request_var', default=None)
+request_contextvar: contextvars.ContextVar[Optional[Request]] = contextvars.ContextVar('request_var', default=None)
 
 
 class ReadOnlyDict(MutableMapping):
@@ -97,16 +97,16 @@ class Storage:
         self.storage_dir = Path('.nicegui')
         self.storage_dir.mkdir(exist_ok=True)
         self._general = PersistentDict(self.storage_dir / 'storage_general.json')
-        self._users = PersistentDict(self.storage_dir / 'storage_users.json')
+        self._users: Dict[str, PersistentDict] = {}
 
     @property
-    def browser(self) -> Dict:
+    def browser(self) -> Union[ReadOnlyDict, Dict]:
         """Small storage that is saved directly within the user's browser (encrypted cookie).
 
         The data is shared between all browser tabs and can only be modified before the initial request has been submitted.
         It is normally better to use `app.storage.user` instead to reduce payload, gain improved security and have larger storage capacity.
         """
-        request: Request = request_contextvar.get()
+        request: Optional[Request] = request_contextvar.get()
         if request is None:
             if globals.get_client() == globals.index_client:
                 raise RuntimeError('app.storage.browser can only be used with page builder functions '
@@ -127,16 +127,17 @@ class Storage:
         The data is stored in a file on the server.
         It is shared between all browser tabs by identifying the user via session cookie ID.
         """
-        request: Request = request_contextvar.get()
+        request: Optional[Request] = request_contextvar.get()
         if request is None:
             if globals.get_client() == globals.index_client:
                 raise RuntimeError('app.storage.user can only be used with page builder functions '
                                    '(https://nicegui.io/documentation/page)')
             else:
                 raise RuntimeError('app.storage.user needs a storage_secret passed in ui.run()')
-        if request.session['id'] not in self._users:
-            self._users[request.session['id']] = {}
-        return self._users[request.session['id']]
+        id = request.session['id']
+        if id not in self._users:
+            self._users[id] = PersistentDict(self.storage_dir / f'storage_user_{id}.json')
+        return self._users[id]
 
     @property
     def general(self) -> Dict:
@@ -145,7 +146,8 @@ class Storage:
 
     async def backup(self) -> None:
         await self._general.backup()
-        await self._users.backup()
+        for user in self._users.values():
+            await user.backup()
 
     async def _loop(self) -> None:
         while True:

+ 3 - 3
tests/test_storage.py

@@ -86,7 +86,7 @@ async def test_access_user_storage_on_interaction(screen: Screen):
     screen.click('switch')
     screen.wait(1)
     await app.storage.backup()
-    assert '{"test_switch": true}' in app.storage._users.filepath.read_text()
+    assert '{"test_switch": true}' in list(app.storage._users.values())[0].filepath.read_text()
 
 
 def test_access_user_storage_from_button_click_handler(screen: Screen):
@@ -102,7 +102,7 @@ def test_access_user_storage_from_button_click_handler(screen: Screen):
     screen.open('/')
     screen.click('test')
     screen.wait(1)
-    assert '{"inner_function": "works"}' in app.storage._users.filepath.read_text()
+    assert '{"inner_function": "works"}' in list(app.storage._users.values())[0].filepath.read_text()
 
 
 async def test_access_user_storage_from_background_task(screen: Screen):
@@ -116,7 +116,7 @@ async def test_access_user_storage_from_background_task(screen: Screen):
 
     screen.ui_run_kwargs['storage_secret'] = 'just a test'
     screen.open('/')
-    assert '{"subtask": "works"}' in app.storage._users.filepath.read_text()
+    assert '{"subtask": "works"}' in list(app.storage._users.values())[0].filepath.read_text()
 
 
 def test_user_and_general_storage_is_persisted(screen: Screen):