2
2
from typing import Awaitable , Callable , List , NoReturn , cast
3
3
4
4
from aiohttp import web
5
+ from aiohttp .typedefs import Handler
5
6
from cryptography import fernet
6
7
7
8
from aiohttp_session import get_session , setup
8
9
from aiohttp_session .cookie_storage import EncryptedCookieStorage
9
10
10
- _Handler = Callable [[web .Request ], Awaitable [web .StreamResponse ]]
11
-
12
11
13
12
def flash (request : web .Request , message : str ) -> None :
14
13
request .setdefault ("flash_outgoing" , []).append (message )
@@ -18,17 +17,16 @@ def get_messages(request: web.Request) -> List[str]:
18
17
return cast (List [str ], request .pop ("flash_incoming" ))
19
18
20
19
21
- async def flash_middleware (app : web .Application , handler : _Handler ) -> _Handler :
22
- async def process (request : web .Request ) -> web .StreamResponse :
23
- session = await get_session (request )
24
- request ["flash_incoming" ] = session .pop ("flash" , [])
25
- response = await handler (request )
20
+ @web .middleware
21
+ async def flash_middleware (request : web .Request , handler : Handler ) -> web .StreamResponse :
22
+ session = await get_session (request )
23
+ request ["flash_incoming" ] = session .pop ("flash" , [])
24
+ try :
25
+ return await handler (request )
26
+ finally :
26
27
session ["flash" ] = request .get ("flash_incoming" , []) + request .get (
27
28
"flash_outgoing" , []
28
29
)
29
- return response
30
-
31
- return process
32
30
33
31
34
32
async def flash_handler (request : web .Request ) -> NoReturn :
@@ -50,10 +48,10 @@ def make_app() -> web.Application:
50
48
fernet_key = fernet .Fernet .generate_key ()
51
49
secret_key = base64 .urlsafe_b64decode (fernet_key )
52
50
setup (app , EncryptedCookieStorage (secret_key ))
51
+ # Install flash middleware (must be installed after aiohttp-session middleware).
52
+ app .middlewares .append (flash_middleware )
53
53
app .router .add_get ("/" , handler )
54
54
app .router .add_get ("/flash" , flash_handler )
55
- # Install flash middleware
56
- app .middlewares .append (flash_middleware )
57
55
return app
58
56
59
57
0 commit comments