add missing type hints to util code

policy_sentry/querying/actions.py

@@ -44,7 +44,7 @@ def get_actions_for_service(service_prefix: str) -> list[str]:
     """
     service_prefix_data = get_service_prefix_data(service_prefix)
     results = []
-    if isinstance(service_prefix_data, dict):
+    if service_prefix_data and isinstance(service_prefix_data, dict):
         results = [
             f"{service_prefix}:{item}" for item in service_prefix_data["privileges"]
         ]
@@ -406,6 +406,9 @@ def get_actions_matching_arn_v2(arn: str) -> list[str]:
     results = set()
     for raw_arn in raw_arns:
         resource_type_name = get_resource_type_name_with_raw_arn(raw_arn)
+        if resource_type_name is None:
+            continue
+
         service_prefix = get_service_from_arn(raw_arn)
         service_prefix_data = get_service_prefix_data(service_prefix)
         for action_name, action_data in service_prefix_data["privileges"].items():
@@ -415,7 +418,7 @@ def get_actions_matching_arn_v2(arn: str) -> list[str]:
     return list(results)
 
 
-def get_actions_matching_condition_key(service_prefix: str, condition_key: str):
+def get_actions_matching_condition_key(service_prefix: str, condition_key: str) -> list[str]:
     """
     Get a list of actions under a service that allow the use of a specified condition key
 
@@ -610,7 +613,7 @@ def get_privilege_info(service_prefix: str, action: str) -> dict[str, Any]:
         List: The info from the docs about that action, along with some of the info from the docs
     """
     try:
-        privilege_info = iam_definition[service_prefix]["privileges"][action]
+        privilege_info: dict[str, Any] = iam_definition[service_prefix]["privileges"][action]
         privilege_info["service_resources"] = iam_definition[service_prefix][
             "resources"
         ]
@@ -637,22 +640,22 @@ def get_api_documentation_link_for_action(
         List: Link to the documentation about that API call
     """
     rows = get_action_data(service_prefix, action_name)
-    for row in rows.get(service_prefix):
-        doc_link = row.get("api_documentation_link")
+    for row in rows.get(service_prefix, []):
+        doc_link: str | None = row.get("api_documentation_link")
         if doc_link:
             return doc_link
     return None
 
 
 @functools.lru_cache(maxsize=1024)
-def get_all_action_links() -> dict[str, str]:
+def get_all_action_links() -> dict[str, str | None]:
     """
     Gets a huge list of the links to all AWS IAM actions. This is meant for use by Cloudsplaining.
 
     :return: A dictionary of all actions present in the database, with the values being the API documentation links.
     """
     all_actions = get_all_actions()
-    results = {}
+    results: dict[str, str | None] = {}
     for action in all_actions:
         try:
             service_prefix, action_name = action.split(":")

policy_sentry/querying/actions_v1.py

@@ -227,6 +227,9 @@ def get_actions_matching_arn_v1(arn: str) -> list[str]:
     results = []
     for raw_arn in raw_arns:
         resource_type_name = get_resource_type_name_with_raw_arn(raw_arn)
+        if resource_type_name is None:
+            continue
+
         service_prefix = get_service_from_arn(raw_arn)
         service_prefix_data = get_service_prefix_data(service_prefix)
         for action_name, action_data in service_prefix_data["privileges"].items():

policy_sentry/querying/all.py

@@ -68,7 +68,7 @@ def get_all_actions_v2(
     :param lowercase: Set to true to have the list of actions be in all lowercase strings.
     :return: A list of all actions present in the database.
     """
-    all_actions = set()
+    all_actions: set[str] = set()
 
     for service_prefix in all_service_prefixes:
         service_prefix_data = get_service_prefix_data(service_prefix)
@@ -88,5 +88,5 @@ def get_service_authorization_url(service_prefix: str) -> str | None:
     """
     Gets the URL to the Actions, Resources, and Condition Keys page for a particular service.
     """
-    result = iam_definition.get(service_prefix, {}).get("service_authorization_url")
+    result: str = iam_definition.get(service_prefix, {}).get("service_authorization_url")
     return result

policy_sentry/querying/arns.py

@@ -7,7 +7,7 @@
 import logging
 import functools
 import warnings
-from typing import Any
+from typing import Any, cast
 
 from policy_sentry.querying.arns_v1 import get_arn_type_details_v1
 from policy_sentry.shared.constants import POLICY_SENTRY_SCHEMA_VERSION_V2
@@ -146,7 +146,7 @@ def get_resource_type_name_with_raw_arn(raw_arn: str) -> str | None:
 
     for resource_name, resource_data in service_data["resources"].items():
         if resource_data["arn"].lower() == raw_arn.lower():
-            return resource_name
+            return cast("str", resource_name)
 
     return None
 

policy_sentry/shared/iam_data.py

@@ -5,7 +5,7 @@
 import logging
 import functools
 from pathlib import Path
-from typing import Any
+from typing import Any, cast
 
 from policy_sentry.shared.constants import (
     DATASTORE_FILE_PATH,
@@ -22,8 +22,11 @@
 
 @functools.lru_cache(maxsize=1)
 def get_iam_definition_schema_version() -> str:
-    return iam_definition.get(
-        POLICY_SENTRY_SCHEMA_VERSION_NAME, POLICY_SENTRY_SCHEMA_VERSION_V1
+    return cast(
+        "str",
+        iam_definition.get(
+            POLICY_SENTRY_SCHEMA_VERSION_NAME, POLICY_SENTRY_SCHEMA_VERSION_V1
+        ),
     )
 
 
@@ -37,10 +40,8 @@ def get_service_prefix_data(service_prefix: str) -> dict[str, Any]:
     Returns:
         List: A list of metadata about that service
     """
-    # result = list(filter(lambda item: item["prefix"] == service_prefix, iam_definition))
-    result = iam_definition.get(service_prefix, None)
     try:
-        return result
+        return cast("dict[str, Any]", iam_definition.get(service_prefix, {}))
     # pylint: disable=bare-except, inconsistent-return-statements
     except:
         logger.info(f"Service prefix not {service_prefix} found.")