Enforce types in function calls

3 files changed, 14 insertions, 7 deletions

diff --git a/accounts/utils/confirmation.py b/accounts/utils/confirmation.py
index b75716d..60967de 100644
--- a/accounts/utils/confirmation.py
+++ b/accounts/utils/confirmation.py
@@ -3,6 +3,8 @@ from itsdangerous import BadSignature, SignatureExpired, URLSafeTimedSerializer
 from werkzeug.exceptions import Forbidden
 from accounts.app import accounts_app
 
+from typing import Union, Optional, Any
+
 
 class Confirmation(URLSafeTimedSerializer):
 
@@ -11,7 +13,10 @@ class Confirmation(URLSafeTimedSerializer):
             key = accounts_app.config['SECRET_KEY']
         super(Confirmation, self).__init__(key, salt=realm, **kwargs)
 
-    def loads_http(self, s, max_age=None, return_timestamp=False, salt=None):
+    def loads_http(self, s: Union[str, bytes],
+                   max_age: Optional[int] = None,
+                   return_timestamp: bool = False,
+                   salt: Optional[bytes] = None) -> Any:
         """
         Like `Confirmation.loads`, but raise HTTP exceptions with appropriate
         messages instead of `BadSignature` or `SignatureExpired`.
diff --git a/accounts/utils/login.py b/accounts/utils/login.py
index 0cd1dc4..64e6ce8 100644
--- a/accounts/utils/login.py
+++ b/accounts/utils/login.py
@@ -7,14 +7,16 @@ import json
 import flask_login.login_manager
 from accounts.app import accounts_app
 
+from typing import Union, Any
+
 
 class _compact_json:
     @staticmethod
-    def loads(payload):
+    def loads(payload: Union[bytes, str, bytearray]) -> Any:
         return json.loads(payload)
 
     @staticmethod
-    def dumps(obj, **kwargs):
+    def dumps(obj: Union[list, dict, tuple], **kwargs):
         kwargs.setdefault("ensure_ascii", False)
         kwargs.setdefault("separators", (",", ":"))
         return json.dumps(obj, **kwargs)
@@ -26,7 +28,7 @@ def create_login_manager() -> flask_login.login_manager.LoginManager:
     login_manager.login_view = 'login.login'
 
     @login_manager.user_loader
-    def load_user(user_id: str):
+    def load_user(user_id: str) -> LoginManager:
         try:
             username, password = parse_userid(user_id)
             return accounts_app.user_backend.auth(username, password)
@@ -37,12 +39,12 @@ def create_login_manager() -> flask_login.login_manager.LoginManager:
     return login_manager
 
 
-def create_userid(username: str, password: str):
+def create_userid(username: str, password: str) -> bytes:
     userid = (username, password)
     return base64_encode(_compact_json.dumps(userid))
 
 
-def parse_userid(value: str):
+def parse_userid(value: str) -> Any:
     return _compact_json.loads(base64_decode(value))
 
 
diff --git a/accounts/utils/sessions.py b/accounts/utils/sessions.py
index 007c928..a452fe1 100644
--- a/accounts/utils/sessions.py
+++ b/accounts/utils/sessions.py
@@ -24,7 +24,7 @@ def _unpad(value: str) -> str:
 
 class EncryptedSerializer(TaggedJSONSerializer):
 
-    def __init__(self):
+    def __init__(self) -> None:
         super(EncryptedSerializer, self).__init__()
         self.block_size = AES.block_size