1212import os
1313import sys
1414import time
15+ import typing as t
1516from datetime import datetime
17+ from typing import TYPE_CHECKING , Any , ClassVar , Optional , Protocol , Type
1618from uuid import uuid4
1719
1820import pytz
2224except ImportError :
2325 import pickle # type: ignore[no-redef]
2426
27+ if TYPE_CHECKING :
28+ from redis import Redis
29+ from flask import Flask
30+ from flask .wrappers import Request
31+
2532from flask .sessions import SessionInterface as FlaskSessionInterface
2633from flask .sessions import SessionMixin
2734from itsdangerous import BadSignature , Signer , want_bytes
@@ -38,11 +45,23 @@ def total_seconds(td):
3845 return td .days * 60 * 60 * 24 + td .seconds
3946
4047
48+ class SessionSerializerType (Protocol ):
49+ def dumps (self , session_content : Any ) -> bytes :
50+ """serialize a session object"""
51+ ...
52+
53+ def loads (self , serialized_session : bytes ) -> Any :
54+ """deserialize a session object"""
55+ ...
56+
57+
4158class ServerSideSession (CallbackDict , SessionMixin ):
4259 """Baseclass for server-side based sessions."""
4360
44- def __init__ (self , initial = None , sid = None , permanent = None ):
45- def on_update (self ):
61+ def __init__ (
62+ self , initial = None , sid : Optional [str ] = None , permanent : Optional [bool ] = None
63+ ):
64+ def on_update (self ) -> None :
4665 self .modified = True
4766
4867 CallbackDict .__init__ (self , initial , on_update )
@@ -89,16 +108,25 @@ class DynamoDBSession(ServerSideSession):
89108
90109
91110class PeeweeSession (ServerSideSession ):
92- def __init__ (self , initial = None , sid = None , permanent = None , ip = None ):
111+ def __init__ (
112+ self ,
113+ initial = None ,
114+ sid : Optional [str ] = None ,
115+ permanent : Optional [bool ] = None ,
116+ ip = None ,
117+ ):
93118 super ().__init__ (initial , sid , permanent )
94119 self .ip = ip
95120
96121
97122class SessionInterface (FlaskSessionInterface ):
98- def _generate_sid (self ):
123+ serializer : ClassVar [SessionSerializerType ]
124+ session_class : ClassVar [Type [ServerSideSession ]]
125+
126+ def _generate_sid (self ) -> str :
99127 return str (uuid4 ())
100128
101- def _get_signer (self , app ) :
129+ def _get_signer (self , app : Flask ) -> Optional [ Signer ] :
102130 if not app .secret_key :
103131 return None
104132 return Signer (app .secret_key , salt = "flask-session" , key_derivation = "hmac" )
@@ -107,7 +135,7 @@ def _get_signer(self, app):
107135class NullSessionInterface (SessionInterface ):
108136 """Used to open a :class:`flask.sessions.NullSession` instance."""
109137
110- def open_session (self , app , request ) :
138+ def open_session (self , app : Flask , request : Request ) -> None :
111139 return None
112140
113141
@@ -123,10 +151,16 @@ class RedisSessionInterface(SessionInterface):
123151 :param permanent: Whether to use permanent session or not.
124152 """
125153
126- serializer = pickle
127- session_class = RedisSession
154+ serializer : ClassVar [ SessionSerializerType ] = t . cast ( SessionSerializerType , pickle )
155+ session_class : ClassVar [ Type [ RedisSession ]] = RedisSession
128156
129- def __init__ (self , redis , key_prefix , use_signer = False , permanent = True ):
157+ def __init__ (
158+ self ,
159+ redis : Optional [Redis ],
160+ key_prefix : str ,
161+ use_signer : bool = False ,
162+ permanent : bool = True ,
163+ ):
130164 if redis is None :
131165 from redis import Redis
132166
@@ -137,7 +171,7 @@ def __init__(self, redis, key_prefix, use_signer=False, permanent=True):
137171 self .permanent = permanent
138172 self .has_same_site_capability = hasattr (self , "get_cookie_samesite" )
139173
140- def open_session (self , app , request ) :
174+ def open_session (self , app : Flask , request : Request ) -> Optional [ SessionMixin ] :
141175 sid = request .cookies .get (app .config ["SESSION_COOKIE_NAME" ])
142176 if not sid :
143177 sid = self ._generate_sid ()
@@ -164,7 +198,9 @@ def open_session(self, app, request):
164198 return self .session_class (sid = sid , permanent = self .permanent )
165199 return self .session_class (sid = sid , permanent = self .permanent )
166200
167- def save_session (self , app , session , response ):
201+ def save_session (
202+ self , app : Flask , session : SessionMixin , response : Response
203+ ) -> None :
168204 if not self .should_set_cookie (app , session ):
169205 return
170206 domain = self .get_cookie_domain (app )
@@ -202,7 +238,9 @@ def save_session(self, app, session, response):
202238 )
203239
204240 if self .use_signer :
205- session_id = self ._get_signer (app ).sign (want_bytes (session .sid ))
241+ session_id = t .cast (Signer , self ._get_signer (app )).sign (
242+ want_bytes (session .sid )
243+ )
206244 else :
207245 session_id = session .sid
208246 response .set_cookie (
0 commit comments