From 9d57fcf13749239490db2cbd356b11df8695a1be Mon Sep 17 00:00:00 2001
From: StefanVDWeide <stefan.vander.weide@hotmail.com>
Date: Mon, 1 May 2023 23:43:03 +0200
Subject: [PATCH] Added better types to routes and helper functions

---
 app/auth/routes.py             | 25 +++++++++++++------------
 app/comments/routes.py         | 13 +++++++------
 app/errors/handlers.py         |  6 +++---
 app/helpers/task_helpers.py    |  1 +
 app/helpers/test_helpers.py    |  2 +-
 app/models.py                  |  6 ++++--
 app/posts/routes.py            | 14 +++++++-------
 app/tasks/long_running_jobs.py | 20 ++++++++++++--------
 app/tasks/routes.py            | 17 +++++++----------
 app/users/routes.py            | 13 +++++++------
 flask_api_template.py          | 12 +++++++-----
 11 files changed, 69 insertions(+), 60 deletions(-)

diff --git a/app/auth/routes.py b/app/auth/routes.py
index 2d535b7..795d1b3 100644
--- a/app/auth/routes.py
+++ b/app/auth/routes.py
@@ -1,4 +1,4 @@
-from flask import request, jsonify
+from flask import Response, request, jsonify
 
 from app import db, jwt
 from app.auth import bp
@@ -23,7 +23,8 @@
 @jwt.token_in_blocklist_loader
 def check_if_token_in_blacklist(jwt_header, jwt_data) -> bool:
     """
-    Helper function for checking if a token is present in the database revoked token table
+    Helper function for checking if a token is present in the database 
+    revoked token table
 
     Parameters
     ----------
@@ -42,7 +43,7 @@ def check_if_token_in_blacklist(jwt_header, jwt_data) -> bool:
 
 
 @bp.post("/register")
-def register() -> str:
+def register() -> tuple[Response, int] | Response:
     """
     Endpoint for adding a new user to the database
 
@@ -52,7 +53,7 @@ def register() -> str:
         A JSON object containing the success message
     """
     try:
-        result = user_schema.load(request.json)
+        result = user_schema.load(request.get_json())
 
     except ValidationError as e:
         return bad_request(e.messages)
@@ -63,7 +64,7 @@ def register() -> str:
     if Users.query.filter_by(email=result["email"]).first():
         return bad_request("Email already in use")
 
-    user = Users(
+    user: Users = Users(
         username=result["username"],
         first_name=result["first_name"],
         last_name=result["last_name"],
@@ -80,7 +81,7 @@ def register() -> str:
 
 
 @bp.post("/login")
-def login() -> str:
+def login() -> tuple[Response, int] | Response:
     """
     Endpoint for authorizing a user and retrieving a JWT
 
@@ -90,7 +91,7 @@ def login() -> str:
         A JSON object containing both the access JWT and the refresh JWT
     """
     try:
-        result = user_schema.load(request.json)
+        result = user_schema.load(request.get_json())
 
     except ValidationError as e:
         return bad_request(e.messages)
@@ -110,7 +111,7 @@ def login() -> str:
 
 @bp.post("/refresh")
 @jwt_required(refresh=True)
-def refresh() -> str:
+def refresh() -> tuple[Response, int]:
     """
     Endpoint in order to retrieve a new access JWT using the refresh JWT.
     A non-fresh access token is returned because the password is not involved in this transaction
@@ -128,7 +129,7 @@ def refresh() -> str:
 
 
 @bp.post("/fresh-login")
-def fresh_login() -> str:
+def fresh_login() -> tuple[Response, int] | Response:
     """
     Endpoint for requesting a new fresh access token
 
@@ -138,7 +139,7 @@ def fresh_login() -> str:
         A JSON object containing
     """
     try:
-        result = user_schema.load(request.json)
+        result = user_schema.load(request.get_json())
 
     except ValidationError as e:
         return bad_request(e.messages)
@@ -156,7 +157,7 @@ def fresh_login() -> str:
 
 @bp.delete("/logout/token")
 @jwt_required()
-def logout_access_token() -> str:
+def logout_access_token() -> tuple[Response, int]:
     """
     Endpoint for revoking the current user"s access token
 
@@ -174,7 +175,7 @@ def logout_access_token() -> str:
 
 @bp.delete("/logout/fresh")
 @jwt_required(refresh=True)
-def logout_refresh_token() -> str:
+def logout_refresh_token() -> tuple[Response, int]:
     """
     Endpoint for revoking the current user"s refresh token
 
diff --git a/app/comments/routes.py b/app/comments/routes.py
index 5580704..ff32fa6 100644
--- a/app/comments/routes.py
+++ b/app/comments/routes.py
@@ -1,4 +1,5 @@
-from flask import request, jsonify
+from typing import Any
+from flask import request, jsonify, Response
 
 from app import db
 from app.comments import bp
@@ -20,7 +21,7 @@
 
 @bp.get("/get/user/comments/post/<int:id>")
 @jwt_required()
-def get_comments_by_post_id(id: int) -> str:
+def get_comments_by_post_id(id: int) -> Response:
     """
     Endpoint for retrieving the user comments associated with a particular post
 
@@ -40,7 +41,7 @@ def get_comments_by_post_id(id: int) -> str:
 
 @bp.post("/post/user/submit/comment")
 @jwt_required()
-def submit_comment() -> str:
+def submit_comment() -> tuple[Response, int] | Response:
     """
     Lets users submit a comment regarding a post
 
@@ -50,7 +51,7 @@ def submit_comment() -> str:
         A JSON object containing a success message
     """
     try:
-        result = comment_deserializing_schema.load(request.json)
+        result = comment_deserializing_schema.load(request.get_json())
     except ValidationError as e:
         return bad_request(e.messages)
 
@@ -72,7 +73,7 @@ def submit_comment() -> str:
 
 @bp.delete("/delete/user/comment/<int:id>")
 @jwt_required()
-def delete_comment(id: int) -> str:
+def delete_comment(id: int) -> tuple[Response, int] | Response:
     """
     Lets users delete one of their own comments
 
@@ -102,7 +103,7 @@ def delete_comment(id: int) -> str:
 
 @bp.get("/get/user/comments/async")
 @jwt_required()
-async def async_comments_api_call() -> str:
+async def async_comments_api_call() -> dict[str, list[Any]]:
     """
     Calls two endpoints from an external API as async demo
 
diff --git a/app/errors/handlers.py b/app/errors/handlers.py
index 09d2718..339870b 100644
--- a/app/errors/handlers.py
+++ b/app/errors/handlers.py
@@ -1,8 +1,8 @@
-from flask import jsonify
+from flask import Response, jsonify
 from werkzeug.http import HTTP_STATUS_CODES
 
 
-def error_response(status_code: int, message=None) -> str:
+def error_response(status_code: int, message=None) -> Response:
     """
     A catch all function which returns the error code and message back to the user
 
@@ -29,7 +29,7 @@ def error_response(status_code: int, message=None) -> str:
     return response
 
 
-def bad_request(message: str) -> str:
+def bad_request(message: str) -> Response:
     """
     Returns a 400 error code when a bad request has been made
 
diff --git a/app/helpers/task_helpers.py b/app/helpers/task_helpers.py
index cd76186..48074da 100644
--- a/app/helpers/task_helpers.py
+++ b/app/helpers/task_helpers.py
@@ -1,4 +1,5 @@
 from rq import get_current_job
+
 from app import db
 from app.models import Tasks
 
diff --git a/app/helpers/test_helpers.py b/app/helpers/test_helpers.py
index 506b028..1ea269d 100644
--- a/app/helpers/test_helpers.py
+++ b/app/helpers/test_helpers.py
@@ -1,4 +1,4 @@
-def register_and_login_test_user(c):
+def register_and_login_test_user(c) -> str:
     """
     Helper function that makes an HTTP request to register a test user
 
diff --git a/app/models.py b/app/models.py
index 60bd716..fa3a1a9 100644
--- a/app/models.py
+++ b/app/models.py
@@ -9,7 +9,8 @@
 @jwt.user_lookup_loader
 def user_loader_callback(jwt_header: dict, jwt_data: dict) -> object:
     """
-    HUser loader function which uses the JWT identity to retrieve a user object. Method is called on protected routes
+    HUser loader function which uses the JWT identity to retrieve a user object. 
+    Method is called on protected routes
 
     Parameters
     ----------
@@ -53,7 +54,8 @@ def set_password(self, password: str):
 
     def check_password(self, password: str) -> bool:
         """
-        Helper function to verify the password hash agains the password provided by the user when logging in
+        Helper function to verify the password hash agains the password provided 
+        by the user when logging in
 
         Parameters
         ----------
diff --git a/app/posts/routes.py b/app/posts/routes.py
index 97a50e2..2d69fac 100644
--- a/app/posts/routes.py
+++ b/app/posts/routes.py
@@ -1,4 +1,4 @@
-from flask import request, jsonify
+from flask import request, jsonify, Response
 
 from app import db
 from app.posts import bp
@@ -20,7 +20,7 @@
 
 @bp.get("get/user/posts")
 @jwt_required()
-def get_posts() -> str:
+def get_posts() -> tuple[Response, int]:
     """
     Returns all posts submitted by the user making the request
 
@@ -36,7 +36,7 @@ def get_posts() -> str:
 
 @bp.get("/get/user/post/<int:id>")
 @jwt_required()
-def get_post_by_id(id: int) -> str:
+def get_post_by_id(id: int) -> tuple[Response, int] | Response:
     """
     Returns a specific post based on the ID in the URL
 
@@ -60,7 +60,7 @@ def get_post_by_id(id: int) -> str:
 
 @bp.post("/post/user/submit/post")
 @jwt_required()
-def submit_post() -> str:
+def submit_post() -> tuple[Response, int] | Response:
     """
     Lets users retrieve a user profile when logged in
 
@@ -72,7 +72,7 @@ def submit_post() -> str:
     try:
         result = post_schema.load(request.json)
     except ValidationError as e:
-        return bad_request(e.messages)
+        return bad_request(e.messages[0])
 
     post = Posts(body=result["body"], user=current_user)
 
@@ -84,7 +84,7 @@ def submit_post() -> str:
 
 @bp.delete("/delete/user/post/<int:id>")
 @jwt_required()
-def delete_post(id: int) -> str:
+def delete_post(id: int) -> tuple[Response, int] | Response:
     """
     Lets users retrieve a user profile when logged in
 
@@ -114,7 +114,7 @@ def delete_post(id: int) -> str:
 
 @bp.get("/get/user/posts/async")
 @jwt_required()
-async def async_posts_api_call() -> str:
+async def async_posts_api_call() -> tuple[dict, int]:
     """
     Calls two endpoints from an external API as async demo
 
diff --git a/app/tasks/long_running_jobs.py b/app/tasks/long_running_jobs.py
index 75244c8..ab183c7 100644
--- a/app/tasks/long_running_jobs.py
+++ b/app/tasks/long_running_jobs.py
@@ -1,7 +1,8 @@
+import sys
 from datetime import time
+
 from app import create_app
 from app.helpers.task_helpers import _set_task_progress
-import sys
 
 # Create the app in order to operate within the context of the app
 app = create_app()
@@ -13,16 +14,19 @@ def count_seconds(**kwargs: int) -> None:
     """
     with app.app_context():
         try:
-            number = kwargs.get("number")
-            _set_task_progress(0)
+            number: int | None = kwargs.get("number")
+            
+            if number:
+                _set_task_progress(0)
 
-            i = 0
+                i = 0
 
-            for i in range(0, number):
-                i += 1
-                time.sleep(1)
-                _set_task_progress(100 * i // number)
+                for i in range(0, number):
+                    i += 1
+                    time.sleep(1)
+                    _set_task_progress(100 * i // number)
 
+        # TODO: Make this a specific except type, no bare except
         except:
             app.logger.error("Unhandled exception", exc_info=sys.exc_info())
 
diff --git a/app/tasks/routes.py b/app/tasks/routes.py
index d35d892..fda0a5b 100644
--- a/app/tasks/routes.py
+++ b/app/tasks/routes.py
@@ -1,20 +1,17 @@
-from app.errors.handlers import bad_request
-from re import T
-from flask import jsonify
+from flask import Response, jsonify
+from flask_jwt_extended import current_user, jwt_required
 
 from app import db
-from app.tasks import bp
+from app.errors.handlers import bad_request
 from app.schemas import TasksSchema
-
-from flask_jwt_extended import jwt_required, current_user
-
+from app.tasks import bp
 
 tasks_schema = TasksSchema(many=True)
 
 
 @bp.get("/background-task/count-seconds/<int:number>")
 @jwt_required()
-def background_worker_count_seconds(number: int) -> str:
+def background_worker_count_seconds(number: int) -> tuple[Response, int] | Response:
     """
     Spawn a background task via RQ to perform a long running task
 
@@ -40,7 +37,7 @@ def background_worker_count_seconds(number: int) -> str:
 
 @bp.get("/get/active-background-tasks")
 @jwt_required()
-def active_background_tasks() -> str:
+def active_background_tasks() -> tuple[Response, int] | str:
     """
     Endpoint to retrieve all the active background tasks
 
@@ -55,7 +52,7 @@ def active_background_tasks() -> str:
 
 @bp.get("/get/finished-background-tasks")
 @jwt_required()
-def finished_background_tasks() -> str:
+def finished_background_tasks() -> tuple[Response, int] | str:
     """
     Endpoint to retrieve the finished background tasks
 
diff --git a/app/users/routes.py b/app/users/routes.py
index 116175c..1f2a7e7 100644
--- a/app/users/routes.py
+++ b/app/users/routes.py
@@ -1,9 +1,10 @@
-from app.users import bp
+from flask import Response
+from flask_jwt_extended import current_user, jwt_required
+
+from app.errors.handlers import bad_request
 from app.models import Users
 from app.schemas import UsersSchema
-from app.errors.handlers import bad_request
-from flask_jwt_extended import jwt_required, current_user
-
+from app.users import bp
 
 # Declare database schemas so they can be returned as JSON objects
 user_schema = UsersSchema(exclude=("email", "password_hash"))
@@ -12,7 +13,7 @@
 
 @bp.get("/get/user/profile")
 @jwt_required()
-def user_page() -> str:
+def user_page() -> tuple[Response, int] | str:
     """
     Let's users retrieve their own user information when logged in
 
@@ -26,7 +27,7 @@ def user_page() -> str:
 
 @bp.get("/get/user/profile/<string:username>")
 @jwt_required()
-def get_user(username: str) -> str:
+def get_user(username: str) -> tuple[Response, int] | Response:
     """
     Lets users retrieve a user profile when logged in
 
diff --git a/flask_api_template.py b/flask_api_template.py
index 3532248..9602ed1 100644
--- a/flask_api_template.py
+++ b/flask_api_template.py
@@ -1,19 +1,21 @@
-from app import create_app
-from app import db
-
 from datetime import datetime
+
 from dateutil.relativedelta import relativedelta
 
+from app import create_app, db
+
 app = create_app()
 
 
 @app.cli.command()
 def remove_old_jwts():
     """
-    Scan the database for JWT tokens in the Revoked Token table older than 5 days and remove them.
+    Scan the database for JWT tokens in the Revoked Token table older than 5 days 
+    and remove them.
     """
 
-    # Import within the function to prevent working outside of application context when calling flask --help
+    # Import within the function to prevent working outside of application context 
+    # when calling flask --help
     from app.models import RevokedTokenModel
 
     delete_date = datetime.utcnow() - relativedelta(days=5)