Fix pyright issues

Add missing type annotations to functions that require them.

Author: tvararu

mavis/reporting/helpers/auth_helper.py

@@ -1,13 +1,15 @@
 from datetime import datetime, timedelta, timezone
 import jwt
-import urllib
+from urllib.parse import quote
+from typing import Any
 
 from flask import (
     request,
     session,
     redirect,
     current_app,
 )
+from flask.sessions import SessionMixin
 
 from functools import wraps
 
@@ -34,7 +36,7 @@ def decorated_function(*args, **kwargs):
 
                     log_user_in(verification_data, session)
                     return redirect(
-                        url_helper.url_without_param(request.full_path, "code")
+                        str(url_helper.url_without_param(request.full_path, "code"))
                     )
                 except KeyError:
                     pass
@@ -44,9 +46,9 @@ def decorated_function(*args, **kwargs):
             return_url = url_helper.externalise_current_url(current_app, request)
             target_url = mavis_helper.mavis_url(
                 current_app,
-                "/start?redirect_uri=" + urllib.parse.quote(return_url),
+                "/start?redirect_uri=" + quote(return_url),
             )
-            return redirect(target_url)
+            return redirect(str(target_url))
 
         return f(*args, **kwargs)
 
@@ -85,7 +87,7 @@ def encode_jwt(payload, current_app=current_app):
     return jwt.encode(payload, secret, algorithm="HS512")
 
 
-def log_user_in(data, session=session):
+def log_user_in(data: dict[str, Any], session: SessionMixin):
     jwt_data = data["jwt_data"]
 
     session["last_visit"] = datetime.now().astimezone(timezone.utc)

mavis/reporting/helpers/mavis_helper.py

@@ -1,7 +1,7 @@
 from http import HTTPStatus
 import requests
 import urllib.parse
-import werkzeug
+from werkzeug.exceptions import Unauthorized
 
 from flask import redirect
 
@@ -49,7 +49,7 @@ def api_call(current_app, session, path, params={}):
     response = get_request(url, headers=headers)
     if response.status_code in [HTTPStatus.UNAUTHORIZED, HTTPStatus.FORBIDDEN]:
         session.clear()
-        raise (werkzeug.exceptions.Unauthorized)
+        raise Unauthorized()
 
     return response
 
@@ -60,7 +60,7 @@ def login_and_return_after(current_app, return_url):
         "/start?redirect_uri=" + urllib.parse.quote_plus(return_url),
     )
     current_app.logger.warning("REDIRECTING TO %s", target_url)
-    return redirect(target_url)
+    return redirect(str(target_url))
 
 
 def get_request(url, headers={}):

mavis/reporting/helpers/url_helper.py

@@ -1,7 +1,7 @@
 from urllib.parse import urlparse, parse_qs, parse_qsl, urlencode, urlunparse, urljoin
 
 
-def url_without_param(url, param):
+def url_without_param(url: str, param: str) -> str:
     parsed_url = urlparse(url)
     query_params_as_dict = parse_qs(parsed_url.query)
     if param in query_params_as_dict:
@@ -16,7 +16,7 @@ def url_without_param(url, param):
         return url
 
 
-def externalise_current_url(current_app, request):
+def externalise_current_url(current_app, request) -> str:
     return urljoin(
         current_app.config["ROOT_URL"] or request.host_url, request.full_path
     )

mavis/reporting/views.py

@@ -51,7 +51,7 @@ def download():
 
 
 @main.errorhandler(404)
-def page_not_found():
+def page_not_found(_error):
     return render_template("errors/404.html"), 404
 
 

tests/helpers/test_auth_helper.py

@@ -1,7 +1,8 @@
 from os import urandom
-
+from typing import cast
 from datetime import datetime, timedelta, timezone
 
+from flask.sessions import SessionMixin
 from mavis.reporting.helpers import auth_helper
 
 
@@ -57,7 +58,7 @@ def test_that_log_user_in_sets_last_visit_to_now(app):
     with app.app_context():
         configure_app(app)
         mock_session = {}
-        auth_helper.log_user_in(mock_user_info(), mock_session)
+        auth_helper.log_user_in(mock_user_info(), cast(SessionMixin, mock_session))
         assert mock_session["last_visit"] is not None
         assert datetime.now().astimezone(timezone.utc) - mock_session[
             "last_visit"
@@ -68,7 +69,7 @@ def test_that_log_user_in_copies_cis2_info_from_the_given_data(app):
     mock_session = {}
     with app.app_context():
         configure_app(app)
-        auth_helper.log_user_in(mock_user_info(), mock_session)
+        auth_helper.log_user_in(mock_user_info(), cast(SessionMixin, mock_session))
         assert mock_session["cis2_info"] == mock_user_info()["jwt_data"]["cis2_info"]
 
 
@@ -78,7 +79,7 @@ def test_that_log_user_in_copies_user_from_the_given_data(app):
 
     with app.app_context():
         configure_app(app)
-        auth_helper.log_user_in(fake_data, mock_session)
+        auth_helper.log_user_in(fake_data, cast(SessionMixin, mock_session))
         assert mock_session["user"] == fake_data["jwt_data"]["user"]
 
 
@@ -89,7 +90,7 @@ def test_that_log_user_in_sets_minimal_jwt(
     fake_data = mock_user_info()
     with app.app_context():
         configure_app(app)
-        auth_helper.log_user_in(fake_data, mock_session)
+        auth_helper.log_user_in(fake_data, cast(SessionMixin, mock_session))
         assert mock_session["jwt"] is not None
         jwt_payload = auth_helper.decode_jwt(mock_session["jwt"])
 

tests/helpers/test_mavis_helper.py

@@ -1,7 +1,7 @@
 import json
 from http import HTTPStatus
 import pytest
-import werkzeug
+from werkzeug.exceptions import Unauthorized
 
 from flask import current_app
 from unittest.mock import patch
@@ -60,7 +60,7 @@ def __init__(self, **kwargs):
         self.json_obj = kwargs.get("json_obj", None)
 
     def json(self):
-        return self.json_obj or json.loads(self.text)
+        return self.json_obj or json.loads(self.text or "{}")
 
 
 def test_that_verify_auth_code_is_called_correctly(
@@ -146,7 +146,7 @@ def test_that_an_unauthorized_api_call_raises_an_exception_and_clears_the_sessio
                 "mavis.reporting.helpers.mavis_helper.get_request",
                 return_value=mock_response,
             ):
-                with pytest.raises(werkzeug.exceptions.Unauthorized):
+                with pytest.raises(Unauthorized):
                     mavis_helper.api_call(
                         app,
                         mock_session,