Merge pull request #33 from AllenInstitute/feature/upgrade-mypy-configs

upgrade mypy and change 'no_site_packages = false' in config

Author: rpmcginty

pyproject.toml

@@ -21,6 +21,7 @@ requires-python = ">=3.9"
 dependencies = [
     "boto3~=1.35",
     "aibs-informatics-core~=0.1",
+    "typing-extensions~=4.15; python_version < '3.11'",
 ]
 
 # For dev dependencies: https://peps.python.org/pep-0735/
@@ -32,8 +33,10 @@ dev = [
 ]
 lint = [
     "boto3-stubs[athena,apigateway,batch,ecr,ecs,efs,essential,fsx,logs,secretsmanager,ses,sns,ssm,sts,stepfunctions]",
-    "mypy~=1.13.0",
+    "mypy~=1.18.0",
     "ruff~=0.9",    
+    "types-pytz",
+    "types-requests",
 ]
 release = [
     "build",
@@ -172,7 +175,7 @@ incremental = false
 # https://mypy.readthedocs.io/en/stable/config_file.html#import-discovery
 ignore_missing_imports = true
 follow_imports = "silent"
-no_site_packages = true
+no_site_packages = false
 
 # Untyped definitions and calls
 # https://mypy.readthedocs.io/en/stable/config_file.html#untyped-definitions-and-calls

src/aibs_informatics_aws_utils/athena.py

@@ -1,27 +1,34 @@
 import logging
+import sys
 import time
 from typing import TYPE_CHECKING, List, Literal, Optional, Tuple
 
+if sys.version_info >= (3, 11):
+    # For Python 3.11+
+    from typing import Unpack
+else:  # pragma: no cover
+    # For Python < 3.11
+    from typing_extensions import Unpack
+
+
 from botocore.exceptions import ClientError
 
 from aibs_informatics_aws_utils.core import AWSService
 from aibs_informatics_aws_utils.exceptions import AWSError
 
 if TYPE_CHECKING:  # pragma: no cover
     from mypy_boto3_athena.type_defs import (
-        GetQueryExecutionInputRequestTypeDef,
         GetQueryExecutionOutputTypeDef,
         QueryExecutionStatusTypeDef,
         QueryExecutionTypeDef,
-        StartQueryExecutionInputRequestTypeDef,
+        StartQueryExecutionInputTypeDef,
         StartQueryExecutionOutputTypeDef,
     )
 else:
-    GetQueryExecutionInputRequestTypeDef = dict
     GetQueryExecutionOutputTypeDef = dict
     QueryExecutionStatusTypeDef = dict
     QueryExecutionTypeDef = dict
-    StartQueryExecutionInputRequestTypeDef = dict
+    StartQueryExecutionInputTypeDef = dict
     StartQueryExecutionOutputTypeDef = dict
 
 
@@ -36,11 +43,11 @@ def start_query_execution(
     query_string: str,
     work_group: Optional[str] = None,
     execution_parameters: Optional[List[str]] = None,
-    **kwargs,
+    **kwargs: Unpack[StartQueryExecutionInputTypeDef],
 ) -> StartQueryExecutionOutputTypeDef:
     athena = get_athena_client()
 
-    request = StartQueryExecutionInputRequestTypeDef(QueryString=query_string)
+    request = StartQueryExecutionInputTypeDef(QueryString=query_string)
     if work_group:
         request["WorkGroup"] = work_group
     if execution_parameters:
@@ -73,8 +80,8 @@ def query_waiter(
         logger.info(f"Query Execution Status: {stats}")
         status = stats["QueryExecution"].get("Status", {})
         state = status.get("State")
-        if state and state in ["SUCCEEDED", "FAILED", "CANCELLED"]:
-            return state, status
+        if state in ["SUCCEEDED", "FAILED", "CANCELLED", "TIMEOUT"]:
+            return state, status  # type: ignore[return-value]
         time.sleep(0.2)  # 200ms
         # Exit if the time waiting exceed the timeout seconds
         if time.time() > start + timeout:

src/aibs_informatics_aws_utils/batch.py

@@ -1,12 +1,12 @@
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Any, Dict, List, Mapping, Optional, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Literal, Mapping, Optional, Union
 
 from aibs_informatics_core.env import ENV_BASE_KEY_ALIAS, EnvBase, get_env_base
 from aibs_informatics_core.models.aws.batch import JobName, ResourceRequirements
 from aibs_informatics_core.utils.decorators import retry
 from aibs_informatics_core.utils.hashing import sha256_hexdigest
 from aibs_informatics_core.utils.logging import get_logger
-from aibs_informatics_core.utils.tools.dicttools import convert_key_case, remove_null_values
+from aibs_informatics_core.utils.tools.dicttools import convert_key_case
 from aibs_informatics_core.utils.tools.strtools import pascalcase
 from botocore.exceptions import ClientError
 
@@ -24,8 +24,8 @@
         HostTypeDef,
         JobDefinitionTypeDef,
         KeyValuePairTypeDef,
+        LinuxParametersTypeDef,
         MountPointTypeDef,
-        RegisterJobDefinitionRequestRequestTypeDef,
         RegisterJobDefinitionResponseTypeDef,
         ResourceRequirementTypeDef,
         RetryStrategyTypeDef,
@@ -42,8 +42,8 @@
     JobDefinitionTypeDef = dict
     DescribeJobsResponseTypeDef = dict
     KeyValuePairTypeDef = dict
+    LinuxParametersTypeDef = dict
     MountPointTypeDef = dict
-    RegisterJobDefinitionRequestRequestTypeDef = dict
     RegisterJobDefinitionResponseTypeDef = dict
     ResourceRequirementTypeDef = dict
     RetryStrategyTypeDef = dict
@@ -61,27 +61,27 @@ def to_volume(
     name: Optional[str],
     efs_volume_configuration: Optional[EFSVolumeConfigurationTypeDef],
 ) -> VolumeTypeDef:
-    return remove_null_values(
-        VolumeTypeDef(
-            host=HostTypeDef(sourcePath=source_path) if source_path else None,
-            name=name,
-            efsVolumeConfiguration=efs_volume_configuration,
-        )
-    )
+    volume_dict = VolumeTypeDef()
+    if source_path:
+        volume_dict["host"] = HostTypeDef(sourcePath=source_path)
+    if name:
+        volume_dict["name"] = name
+    if efs_volume_configuration:
+        volume_dict["efsVolumeConfiguration"] = efs_volume_configuration
+    return volume_dict
 
 
 def to_mount_point(
     container_path: Optional[str],
     read_only: bool,
     source_volume: Optional[str],
 ) -> MountPointTypeDef:
-    return remove_null_values(
-        MountPointTypeDef(
-            containerPath=container_path,
-            readOnly=read_only,
-            sourceVolume=source_volume,
-        )
-    )
+    mount_point_dict = MountPointTypeDef(readOnly=read_only)
+    if container_path:
+        mount_point_dict["containerPath"] = container_path
+    if source_volume:
+        mount_point_dict["sourceVolume"] = source_volume
+    return mount_point_dict
 
 
 def to_key_value_pairs(
@@ -104,7 +104,7 @@ def to_key_value_pairs(
             for k, v in environment.items()
             if not remove_null_values or v is not None
         ],
-        key=lambda _: _.get("name"),
+        key=lambda _: _.get("name", ""),
     )
 
 
@@ -115,16 +115,23 @@ def to_resource_requirements(
 ) -> List[ResourceRequirementTypeDef]:
     """Converts Batch resource requirement parameters into a list of ResourceRequirement objects
 
+    The returned list only includes dictionary entries for resources that specify
+    an explicit value. Anything unset will be dropped.
+
     Args:
-        gpu (Optional[int], optional): number of . Defaults to None.
-        memory (Optional[int], optional): _description_. Defaults to None.
-        vcpus (Optional[int], optional): _description_. Defaults to None.
+        gpu (Optional[int], optional): number of GPUs to use. Defaults to None.
+        memory (Optional[int], optional): amount of memory in MiB. Defaults to None.
+        vcpus (Optional[int], optional): Number of VCPUs to use. Defaults to None.
 
     Returns:
         List[ResourceRequirementTypeDef]: list of resource requirements
     """
 
-    pairs = [("GPU", gpu), ("MEMORY", memory), ("VCPU", vcpus)]
+    pairs: list[tuple[Literal["GPU", "MEMORY", "VCPU"], Optional[int]]] = [
+        ("GPU", gpu),
+        ("MEMORY", memory),
+        ("VCPU", vcpus),
+    ]
     return [ResourceRequirementTypeDef(type=t, value=str(v)) for t, v in pairs if v is not None]
 
 
@@ -176,7 +183,7 @@ def register_job_definition(
     tags: Optional[Mapping[str, str]] = None,
     propagate_tags: bool = False,
     region: Optional[str] = None,
-) -> JobDefinitionTypeDef:
+) -> JobDefinitionTypeDef | RegisterJobDefinitionResponseTypeDef:
     batch = get_batch_client(region=region)
 
     # First we check to make sure that we aren't crearting unnecessary revisions
@@ -210,7 +217,7 @@ def register_job_definition(
     logger.info(
         f"Registering job definition with following properties: {register_job_definition_kwargs}"
     )
-    response = batch.register_job_definition(**register_job_definition_kwargs)
+    response = batch.register_job_definition(**register_job_definition_kwargs)  # type: ignore[arg-type]
     return response
 
 
@@ -263,7 +270,7 @@ class BatchJobBuilder:
     mount_points: List[MountPointTypeDef] = field(default_factory=list)
     volumes: List[VolumeTypeDef] = field(default_factory=list)
     privileged: bool = field(default=False)
-    linux_parameters: Optional[Dict[str, Any]] = field(default=None)
+    linux_parameters: Optional[LinuxParametersTypeDef] = field(default=None)
     env_base: EnvBase = field(default_factory=EnvBase.from_env)
 
     def __post_init__(self):
@@ -295,7 +302,7 @@ def container_overrides(self) -> ContainerOverridesTypeDef:
 
     @property
     def container_overrides__sfn(self) -> Dict[str, Any]:
-        return convert_key_case(self.container_overrides, pascalcase)
+        return convert_key_case(self.container_overrides, pascalcase)  # type: ignore[arg-type]
 
     def _normalized_resource_requirements(self) -> List[ResourceRequirementTypeDef]:
         if isinstance(self.resource_requirements, list):

src/aibs_informatics_aws_utils/core.py

@@ -145,7 +145,7 @@ def get_region(region: Optional[str] = None) -> str:
 
 def get_account_id() -> str:
     """Will get the account id from the current credentials/identity"""
-    return get_caller_identity().get("Account")
+    return get_caller_identity()["Account"]
 
 
 def get_user_id() -> UserId:

src/aibs_informatics_aws_utils/data_sync/operations.py

@@ -256,7 +256,7 @@ def sync_s3_to_s3(
         result = DataSyncResult()
         if self.config.include_detailed_response:
             path_stats = get_s3_path_stats(destination_path)
-            result.files_transferred = path_stats.object_count
+            result.files_transferred = path_stats.object_count or 0
             result.bytes_transferred = path_stats.size_bytes
         return result
 

src/aibs_informatics_aws_utils/dynamodb/conditions.py

@@ -255,7 +255,7 @@ def deserialize_condition(
         def _deserialize_condition(ce: ConditionBaseExpression) -> ConditionBase:
             ce_key = (ce.format, ce.operator)
             condition_base_cls = cls._CONDITION_BASE_CLASS_LOOKUP[ce_key]
-            ce_values = []
+            ce_values: list[AttributeBase | ConditionBase] = []
             for ce_value in ce.values:
                 if isinstance(ce_value, ConditionBaseExpression):
                     ce_values.append(_deserialize_condition(ce_value))

src/aibs_informatics_aws_utils/dynamodb/functions.py

@@ -14,18 +14,21 @@
 
 if TYPE_CHECKING:  # pragma: no cover
     from mypy_boto3_dynamodb.type_defs import (
-        BatchGetItemInputRequestTypeDef,
-        GetItemInputRequestTypeDef,
+        BatchGetItemInputTypeDef,
+        GetItemInputTypeDef,
         KeysAndAttributesTypeDef,
-        QueryInputRequestTypeDef,
-        ScanInputRequestTypeDef,
+        QueryInputTableQueryTypeDef,
+        QueryInputTypeDef,
+        ScanInputTypeDef,
     )
 else:
-    BatchGetItemInputRequestTypeDef = dict
+    BatchGetItemInputTypeDef = dict
     GetItemInputRequestTypeDef = dict
+    GetItemInputTypeDef = dict
     KeysAndAttributesTypeDef = dict
-    QueryInputRequestTypeDef = dict
-    ScanInputRequestTypeDef = dict
+    QueryInputTableQueryTypeDef = dict
+    QueryInputTypeDef = dict
+    ScanInputTypeDef = dict
 
 
 logger = get_logger(__name__)
@@ -105,7 +108,7 @@ def table_get_items(
                 "Keys": serialized_keys,
             }
         }
-        props: BatchGetItemInputRequestTypeDef = {
+        props: BatchGetItemInputTypeDef = {
             "RequestItems": request_items,
             "ReturnConsumedCapacity": "NONE",
         }
@@ -213,7 +216,7 @@ def table_query(
         key_expr_component.expression_attribute_values__serialized
     )
 
-    db_request: QueryInputRequestTypeDef = {
+    db_request: QueryInputTypeDef = {
         "TableName": table.name,
         "KeyConditionExpression": key_expr_component.condition_expression,
     }
@@ -287,7 +290,7 @@ def table_scan(
     db = get_dynamodb_client(region=region)
     table = table_as_resource(table_name)
 
-    db_request: ScanInputRequestTypeDef = {"TableName": table.name}
+    db_request: ScanInputTypeDef = {"TableName": table.name}
 
     # Handle when filter_expression is provided
     if filter_expression is not None:

src/aibs_informatics_aws_utils/dynamodb/table.py

@@ -204,12 +204,12 @@ def build_optimized_condition_expression_set(
                 new_condition = Key(k).eq(v)
                 if (
                     k in candidate_conditions
-                    and candidate_conditions[k]._values[1:] != new_condition._values[1:]  # type: ignore[union-attr]
+                    and candidate_conditions[k]._values[1:] != new_condition._values[1:]  # type: ignore[attr-defined,union-attr]
                 ):
                     raise DBQueryException(f"Multiple values provided for attribute {k}!")
                 candidate_conditions[k] = Key(k).eq(v)
-        elif len(_._values) and isinstance(_._values[0], (Key, Attr)):  # type: ignore[union-attr]
-            attr_name = cast(str, _._values[0].name)  # type: ignore[union-attr]
+        elif len(_._values) and isinstance(_._values[0], (Key, Attr)):  # type: ignore[attr-defined,union-attr]
+            attr_name = cast(str, _._values[0].name)  # type: ignore[attr-defined,union-attr]
             if attr_name not in index_all_key_names or not isinstance(
                 _, SupportedKeyComparisonTypes
             ):
@@ -728,7 +728,7 @@ def delete(
         try:
             deleted_attributes = table_delete_item(
                 table_name=self.table_name,
-                key=key,
+                key=cast(DynamoDBKey, key),
                 return_values="ALL_OLD",  # type: ignore[arg-type] # expected type more general than specified here
             )