Fix MyPy errors

status_function/pyproject.toml

@@ -32,6 +32,7 @@ pulumi-azure = "^4.18.0"
 pylint-absolute-imports = "^1.0.1"
 isort = "^5.12.0"
 mypy = "^1.8.0"
+types-requests = "^2.31.0.20240125"
 
 [build-system]
 requires = ["poetry-core"]

status_function/run_tests.sh

@@ -38,6 +38,11 @@ status=$((status+$?))
 python -m coverage report --show-missing
 status=$((status+$?))
 
+# Check types with MyPy
+echo "Running type checking..."
+python -m mypy --config-file tests/mypy.ini status/ tests/
+status=$((status+$?))
+
 # [optional] Check Markdown coding style with Ruby's markdown lint
 # https://github.com/markdownlint/markdownlint
 if ! command -v mdl &> /dev/null

status_function/status/__init__.py

@@ -10,6 +10,7 @@
 from datetime import datetime
 from functools import lru_cache
 from typing import Any, Optional
+from uuid import UUID
 
 import azure.functions as func
 import requests
@@ -24,6 +25,7 @@
 from status import models, settings
 from status.auth import BearerAuth
 from status.logutils import set_log_handler
+from status.models import RoleAssignment
 from status.wrapper import CredentialWrapper
 
 logging.basicConfig(
@@ -241,7 +243,7 @@ def get_role_assignment_models(
 def get_subscription_role_assignment_models(
     subscription: Any,
     graph_client: GraphRbacManagementClient,
-) -> list:
+) -> list[RoleAssignment]:
     """Get the role assignment models for each a subscription.
 
     Args:
@@ -259,7 +261,7 @@ def get_subscription_role_assignment_models(
         for assignment in assignments_list:
             role_assignments_models += get_role_assignment_models(
                 assignment,
-                role_def_dict.get(assignment.properties.role_definition_id),
+                role_def_dict.get(assignment.properties.role_definition_id, "Unknown"),
                 graph_client,
             )
     except CloudError as e:
@@ -271,7 +273,7 @@ def get_subscription_role_assignment_models(
     return role_assignments_models
 
 
-def get_all_status(tenant_id: str) -> list[models.SubscriptionStatus]:
+def get_all_status(tenant_id: UUID) -> list[models.SubscriptionStatus]:
     """Get status and role assignments for all subscriptions.
 
     Args:
@@ -284,7 +286,7 @@ def get_all_status(tenant_id: str) -> list[models.SubscriptionStatus]:
     started_at = datetime.now()
 
     graph_client = GraphRbacManagementClient(
-        credentials=GRAPH_CREDENTIALS, tenant_id=tenant_id
+        credentials=GRAPH_CREDENTIALS, tenant_id=str(tenant_id)
     )
 
     client = SubscriptionClient(credential=CREDENTIALS)
@@ -303,7 +305,7 @@ def get_all_status(tenant_id: str) -> list[models.SubscriptionStatus]:
                 subscription_id=subscription.subscription_id,
                 display_name=subscription.display_name,
                 state=subscription.state,
-                role_assignments=role_assignments_models,
+                role_assignments=tuple(role_assignments_models),
             )
         )
 

status_function/status/auth.py

@@ -6,10 +6,15 @@
         value is 5.
 """
 from datetime import datetime, timedelta
+from typing import Any
 
 import jwt
 import requests
 from cryptography.hazmat.primitives import serialization
+from cryptography.hazmat.primitives.asymmetric import dsa
+from cryptography.hazmat.primitives.asymmetric.ec import EllipticCurvePrivateKey
+from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
+from cryptography.hazmat.primitives.asymmetric.rsa import RSAPrivateKey
 
 from status.settings import get_settings
 
@@ -32,17 +37,21 @@ def __init__(self) -> None:
         # Generate keys with ssh-keygen -t rsa
         private_key_txt = settings.PRIVATE_KEY
 
-        self.private_key = serialization.load_ssh_private_key(
+        private_key = serialization.load_ssh_private_key(
             private_key_txt.encode(), password=b""
         )
+        assert not isinstance(
+            self.private_key, dsa.DSAPrivateKey
+        ), "DSA keys are not supported."
+        self.private_key: EllipticCurvePrivateKey | RSAPrivateKey | Ed25519PrivateKey = private_key  # type: ignore
 
     def create_access_token(self) -> str:
         """Create an access token for the user to access the API.
 
         Returns:
             The access token.
         """
-        token_claims = {"sub": "status-app"}
+        token_claims: dict[str, Any] = {"sub": "status-app"}
         access_token_expires = timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
 
         expire = datetime.utcnow() + access_token_expires

status_function/status/logutils.py

@@ -1,5 +1,6 @@
 """Utils to send logs to Azure Application Insights."""
 import logging
+from typing import Optional
 
 from opencensus.ext.azure.log_exporter import AzureLogHandler
 
@@ -9,7 +10,7 @@
 class CustomDimensionsFilter(logging.Filter):
     """Add application-wide properties to AzureLogHandler records."""
 
-    def __init__(self, custom_dimensions: dict = None) -> None:
+    def __init__(self, custom_dimensions: Optional[dict] = None) -> None:
         """Add custom dimensions to the log record.
 
         Args:

status_function/status/settings.py

@@ -36,9 +36,7 @@ class Config:
         env_file_encoding = "utf-8"
 
     @validator("PRIVATE_KEY")
-    def correct_start_and_end(
-        cls: Config, v: str
-    ) -> str:  # pylint: disable=no-self-argument
+    def correct_start_and_end(cls, v: str) -> str:  # pylint: disable=no-self-argument
         """Check that the private key is a private key.
 
         Args:
@@ -69,4 +67,4 @@ def get_settings() -> Settings:
     Returns:
         The settings.
     """
-    return Settings()
+    return Settings()  # type: ignore

status_function/status/wrapper.py

@@ -1,17 +1,20 @@
 """Copied from https://stackoverflow.com/a/64129363/3324095 and slightly modified."""
+from typing import Optional
+
 from azure.core.pipeline import PipelineContext, PipelineRequest
 from azure.core.pipeline.policies import BearerTokenCredentialPolicy
 from azure.core.pipeline.transport import HttpRequest
 from azure.identity import DefaultAzureCredential
 from msrest.authentication import BasicTokenAuthentication
+from requests import Session
 
 
 class CredentialWrapper(BasicTokenAuthentication):
     """Wrapper for azure-identity credentials."""
 
     def __init__(
         self,
-        credential: DefaultAzureCredential = None,
+        credential: Optional[DefaultAzureCredential] = None,
         resource_id: str = "https://management.azure.com/.default",
         **kwargs: dict
     ):
@@ -29,7 +32,7 @@ def __init__(
         Keyword Args:
             Any other parameter accepted by BasicTokenAuthentication
         """
-        super().__init__(None)
+        self.scheme = "Bearer"
         if credential is None:
             credential = DefaultAzureCredential(
                 exclude_visual_studio_code_credential=True
@@ -58,7 +61,7 @@ def set_token(self):
         token = request.http_request.headers["Authorization"].split(" ", 1)[1]
         self.token = {"access_token": token}
 
-    def signed_session(self, session=None) -> DefaultAzureCredential:
+    def signed_session(self, session=None) -> Session:
         """Sign the session.
 
         Args:

status_function/tests/mypy.ini

@@ -0,0 +1,2 @@
+[mypy]
+ignore_missing_imports = True

usage_function/pyproject.toml

@@ -33,6 +33,7 @@ pylint = "^2.10.2"
 [tool.poetry.group.dev.dependencies]
 isort = "^5.12.0"
 pylint-absolute-imports = "^1.0.1"
+mypy = "^1.8.0"
 
 [tool.black]
 line-width = 88

usage_function/run_tests.sh

@@ -40,6 +40,11 @@ status=$((status+$?))
 python -m coverage report --show-missing
 status=$((status+$?))
 
+# Check types with MyPy
+echo "Running type checking..."
+python -m mypy --config-file tests/mypy.ini usage/ tests/
+status=$((status+$?))
+
 # [optional] Check Markdown coding style with Ruby's markdown lint
 # https://github.com/markdownlint/markdownlint
 if ! command -v mdl &> /dev/null