Fix Pylance type warnings across codebase

Author: jgbustos

ml_rest_api/api/model/predict.py

@@ -1,10 +1,10 @@
 """This module implements the ModelPredict class."""
 
 import json
-from typing import Type, Dict
-from aniso8601 import parse_date, parse_datetime
+from typing import Any
+from aniso8601 import parse_date, parse_datetime  # type: ignore
 from flask import request
-from flask_restx import Resource, Model, fields
+from flask_restx import Resource, Model, fields  # type: ignore
 from ml_rest_api.api.restx import api, FlaskApiReturnType, MLRestAPINotReadyException
 from ml_rest_api.ml_trained_model.wrapper import trained_model_wrapper
 
@@ -14,19 +14,19 @@ def build_api_model() -> Model:
     Returns a Flask-RESTX Api Model based on the sample dict returned by the trained model wrapper.
     This will be used to validate input and automatically generate the Swagger prototype.
     """
-    fields_classes_map: Dict = {
+    fields_classes_map: dict[str, Any] = {
         "str": fields.String,
         "int": fields.Integer,
         "float": fields.Float,
         "bool": fields.Boolean,
         "datetime": fields.DateTime,
         "date": fields.Date,
     }
-    model_dict: Dict = {}
-    model_sample: Dict = trained_model_wrapper.sample()
+    model_dict: dict[str, Any] = {}
+    model_sample: dict[str, Any] = trained_model_wrapper.sample()
     if model_sample:
         for key, value in model_sample.items():
-            fields_class: Type[fields.Raw] = fields_classes_map.get(
+            fields_class: Any = fields_classes_map.get(
                 type(value).__name__, fields.String
             )
             if type(value).__name__ == "str":
@@ -44,14 +44,14 @@ def build_api_model() -> Model:
     return api.model("input_vector", model_dict)
 
 
-ns = api.namespace(  # pylint: disable=invalid-name
+ns = api.namespace(  # pylint: disable=invalid-name  # type: ignore[misc]
     "model",
     description="Methods supported by our ML model",
     validate=bool(trained_model_wrapper.sample()),
 )
 
 
-@ns.route("/predict")
+@ns.route("/predict")  # type: ignore[misc]
 class ModelPredict(Resource):
     """Implements the /model/predict POST method."""
 

ml_rest_api/api/restx.py

@@ -1,10 +1,10 @@
 """Module that creates the Api object and declares default error handler."""
 
 from logging import Logger, getLogger
-from typing import Tuple, Dict
+from typing import Any
 from jsonschema import FormatChecker
 from flask import Blueprint
-from flask_restx import Api
+from flask_restx import Api  # type: ignore
 from ml_rest_api.settings import get_value
 
 
@@ -16,7 +16,7 @@ class MLRestAPINotReadyException(MLRestAPIException):
     """Base ML Rest API NOT READY Exception"""
 
 
-FlaskApiReturnType = Tuple[Dict, int]
+FlaskApiReturnType = tuple[dict[str, Any], int]
 
 log: Logger = getLogger(__name__)
 
@@ -39,15 +39,15 @@ class MLRestAPINotReadyException(MLRestAPIException):
 )
 
 
-@api.errorhandler(MLRestAPINotReadyException)
+@api.errorhandler(MLRestAPINotReadyException)  # type: ignore[misc]
 def not_ready_error_handler() -> FlaskApiReturnType:
     """NOT READY error handler that returns HTTP 503 error."""
     log.exception("Server Not Ready")
     return {"message": "Server Not Ready"}, 503
 
 
-@api.errorhandler
-def default_error_handler(exception) -> FlaskApiReturnType:
+@api.errorhandler  # type: ignore[misc]
+def default_error_handler(exception: Any) -> FlaskApiReturnType:
     """Default error handler that returns HTTP 500 error."""
     log.exception(exception.message)
     if get_value("FLASK_DEBUG"):

ml_rest_api/app.py

@@ -4,27 +4,25 @@
 import warnings
 from logging import Logger, getLogger
 import logging.config
-from typing import List
 from flask import Flask
-from flask_wtf import CSRFProtect  # pylint: disable=unused-import
+from flask_wtf import CSRFProtect  # pylint: disable=unused-import  # type: ignore
 from ml_rest_api.settings import get_value
 from ml_rest_api.ml_trained_model.wrapper import trained_model_wrapper
 from ml_rest_api.api.restx import blueprint
-import ml_rest_api.api.health.liveness  # pylint: disable=unused-import
-import ml_rest_api.api.health.readiness  # pylint: disable=unused-import
-import ml_rest_api.api.model.predict  # pylint: disable=unused-import
+import ml_rest_api.api.health.liveness  # pylint: disable=unused-import  # pyright: ignore[reportUnusedImport]
+import ml_rest_api.api.health.readiness  # pylint: disable=unused-import  # pyright: ignore[reportUnusedImport]
+import ml_rest_api.api.model.predict  # pylint: disable=unused-import  # pyright: ignore[reportUnusedImport]
 
 IN_UWSGI: bool = True  # pylint: disable=invalid-name
 try:
-    # pyright: reportMissingImports=false
-    import uwsgi  # pylint: disable=unused-import
+    import uwsgi  # pylint: disable=unused-import  # type: ignore
 except ImportError:
     IN_UWSGI = False  # pylint: disable=invalid-name
 
 
 def configure_app(flask_app: Flask) -> None:
     """Configures the app."""
-    flask_settings_to_apply: List = [
+    flask_settings_to_apply: list[str] = [
         #'FLASK_SERVER_NAME',
         "SWAGGER_UI_DOC_EXPANSION",
         "RESTX_VALIDATE",

ml_rest_api/ml_trained_model/wrapper.py

@@ -5,10 +5,10 @@
 import importlib
 from threading import Thread
 from types import ModuleType
-from typing import Optional, Iterable, Callable, Dict
+from typing import Any, Optional, Iterable, Callable
 from ml_rest_api.settings import get_value
 
-WrapperCallableType = Optional[Callable]
+WrapperCallableType = Optional[Callable[..., Any]]
 
 
 class TrainedModelWrapper:
@@ -85,13 +85,13 @@ def init(self) -> None:
             self._init()
             self.initialised = True
 
-    def run(self, data: Iterable) -> Dict:
+    def run(self, data: Iterable[Any]) -> dict[str, Any]:
         """Calls the wrapped run() method if it's assigned."""
         if self._run:
             return self._run(data)
         return {}
 
-    def sample(self) -> Dict:
+    def sample(self) -> dict[str, Any]:
         """Calls the wrapped sample() method if it's assigned."""
         if self._sample:
             return self._sample()

ml_rest_api/settings.py

@@ -1,12 +1,12 @@
 """Settings file."""
 
 import os
-from typing import Any, Dict
+from typing import Any
 
 
 def get_value(key: str) -> Any:
     """Returns a value from the corresponding env var or from settings if env var doesn't exist."""
-    settings: Dict = {
+    settings: dict[str, Any] = {
         # Flask settings
         "FLASK_SERVER_NAME": "localhost:8888",
         "FLASK_HOST": "0.0.0.0",

tests/basic_test.py

@@ -2,6 +2,7 @@
 
 from http import HTTPStatus
 from json import loads
+from typing import Any
 import requests
 import pytest
 from openapi_spec_validator import (
@@ -20,7 +21,7 @@
     "Content-Type": "application/json",
     "Accept": "application/json",
 }
-GOOD_JSON_DICT = {
+GOOD_JSON_DICT: dict[str, Any] = {
     "int_param": 12345,
     "string_param": "foobar",
     "float_param": 123.45,
@@ -34,7 +35,7 @@
 NOT_A_DATE_MSG = "is not a 'date'"
 
 
-def _get_request(url):
+def _get_request(url: str):
     response = requests.get(
         url,
         allow_redirects=True,
@@ -43,7 +44,7 @@ def _get_request(url):
     return response
 
 
-def _post_request(url, headers=None, json=None):
+def _post_request(url: str, headers: dict[str, str] | None = None, json: Any = None):
     response = requests.post(
         url,
         headers=headers,
@@ -54,7 +55,9 @@ def _post_request(url, headers=None, json=None):
     return response
 
 
-def _post_request_good_json_with_overrides(override_key=None, override_value=None):
+def _post_request_good_json_with_overrides(
+    override_key: str | None = None, override_value: Any = None
+):
     json_dict = GOOD_JSON_DICT
     if override_key is not None:
         if override_value is None: