♻️Maintenance: mypy fixes for AWS-library, Dask-task-lib (#6116)

Author: sanderegg

.github/workflows/ci-testing-deploy.yml

@@ -899,7 +899,6 @@ jobs:
         run: ./ci/github/unit-testing/aws-library.bash install
       - name: typecheck
         run: ./ci/github/unit-testing/aws-library.bash typecheck
-        continue-on-error: true
       - name: test
         if: always()
         run: ./ci/github/unit-testing/aws-library.bash test
@@ -942,7 +941,6 @@ jobs:
         run: ./ci/github/unit-testing/dask-task-models-library.bash install
       - name: typecheck
         run: ./ci/github/unit-testing/dask-task-models-library.bash typecheck
-        continue-on-error: true
       - name: test
         if: always()
         run: ./ci/github/unit-testing/dask-task-models-library.bash test

packages/aws-library/requirements/_test.txt

@@ -18,13 +18,21 @@ boto3==1.34.131
     #   -c requirements/_base.txt
     #   aws-sam-translator
     #   moto
+boto3-stubs==1.34.149
+    # via types-boto3
 botocore==1.34.131
     # via
     #   -c requirements/_base.txt
     #   aws-xray-sdk
     #   boto3
     #   moto
     #   s3transfer
+botocore-stubs==1.34.144
+    # via
+    #   -c requirements/_base.txt
+    #   boto3-stubs
+    #   types-aioboto3
+    #   types-aiobotocore
 certifi==2024.7.4
     # via
     #   -c requirements/../../../requirements/constraints.txt
@@ -242,15 +250,34 @@ tomli==2.0.1
     # via
     #   coverage
     #   pytest
+types-aioboto3==13.1.1
+    # via -r requirements/_test.in
+types-aiobotocore==2.13.1
+    # via
+    #   -c requirements/_base.txt
+    #   types-aioboto3
+types-awscrt==0.21.2
+    # via
+    #   -c requirements/_base.txt
+    #   botocore-stubs
+types-boto3==1.0.2
+    # via -r requirements/_test.in
+types-s3transfer==0.10.1
+    # via
+    #   boto3-stubs
+    #   types-aioboto3
 typing-extensions==4.12.2
     # via
     #   -c requirements/_base.txt
     #   aws-sam-translator
+    #   boto3-stubs
     #   cfn-lint
     #   flexcache
     #   flexparser
     #   pint
     #   pydantic
+    #   types-aioboto3
+    #   types-aiobotocore
 urllib3==2.2.2
     # via
     #   -c requirements/../../../requirements/constraints.txt

packages/aws-library/setup.cfg

@@ -15,3 +15,7 @@ test = pytest
 
 [tool:pytest]
 asyncio_mode = auto
+
+[mypy]
+plugins =
+	pydantic.mypy

packages/dask-task-models-library/src/dask_task_models_library/container_tasks/events.py

@@ -2,6 +2,7 @@
 from abc import ABC, abstractmethod
 from typing import Any, ClassVar, TypeAlias
 
+import dask.typing
 from distributed.worker import get_worker
 from pydantic import BaseModel, Extra, validator
 
@@ -22,6 +23,14 @@ class Config:
         extra = Extra.forbid
 
 
+def _dask_key_to_dask_task_id(key: dask.typing.Key) -> str:
+    if isinstance(key, bytes):
+        return key.decode("utf-8")
+    if isinstance(key, tuple):
+        return "(" + ", ".join(_dask_key_to_dask_task_id(k) for k in key) + ")"
+    return f"{key}"
+
+
 class TaskProgressEvent(BaseTaskEvent):
     progress: float
 
@@ -37,7 +46,7 @@ def from_dask_worker(
         job_id = worker.get_current_task()
 
         return cls(
-            job_id=job_id,
+            job_id=_dask_key_to_dask_task_id(job_id),
             progress=progress,
             task_owner=task_owner,
         )
@@ -97,7 +106,7 @@ def from_dask_worker(
         worker = get_worker()
         job_id = worker.get_current_task()
         return cls(
-            job_id=job_id,
+            job_id=_dask_key_to_dask_task_id(job_id),
             log=log,
             log_level=log_level,
             task_owner=task_owner,

packages/dask-task-models-library/src/dask_task_models_library/container_tasks/io.py

@@ -1,21 +1,21 @@
 import json
 from contextlib import suppress
 from pathlib import Path
-from typing import Any, ClassVar, Union, cast
+from typing import Any, ClassVar, TypeAlias, Union
 
 from models_library.basic_regex import MIME_TYPE_RE, PROPERTY_KEY_RE
 from models_library.generics import DictModel
 from pydantic import (
     AnyUrl,
     BaseModel,
+    ConstrainedStr,
     Extra,
     Field,
     StrictBool,
     StrictFloat,
     StrictInt,
     StrictStr,
 )
-from pydantic.types import constr
 
 TaskCancelEventName = "cancel_event_{}"
 
@@ -81,8 +81,11 @@ class Config:
         }
 
 
-PortKey = constr(regex=PROPERTY_KEY_RE)
-PortValue = Union[
+class PortKey(ConstrainedStr):
+    regex = PROPERTY_KEY_RE
+
+
+PortValue: TypeAlias = Union[
     StrictBool,
     StrictInt,
     StrictFloat,
@@ -95,7 +98,7 @@ class Config:
 
 
 class TaskInputData(DictModel[PortKey, PortValue]):
-    class Config(DictModel.Config):
+    class Config:
         schema_extra: ClassVar[dict[str, Any]] = {
             "examples": [
                 {
@@ -109,7 +112,7 @@ class Config(DictModel.Config):
         }
 
 
-PortSchemaValue = Union[PortSchema, FilePortSchema]
+PortSchemaValue: TypeAlias = Union[PortSchema, FilePortSchema]
 
 
 class TaskOutputDataSchema(DictModel[PortKey, PortSchemaValue]):
@@ -119,7 +122,7 @@ class TaskOutputDataSchema(DictModel[PortKey, PortSchemaValue]):
     # does not work well in that case. For that reason, the schema is
     # sent as a json-schema instead of with a dynamically-created model class
     #
-    class Config(DictModel.Config):
+    class Config:
         schema_extra: ClassVar[dict[str, Any]] = {
             "examples": [
                 {
@@ -171,10 +174,9 @@ def from_task_output(
                 msg = f"Could not locate '{output_key}' in {output_data_file}"
                 raise ValueError(msg)
 
-        # NOTE: this cast is necessary to make mypy happy
-        return cast(TaskOutputData, cls.parse_obj(data))
+        return cls.parse_obj(data)
 
-    class Config(DictModel.Config):
+    class Config:
         schema_extra: ClassVar[dict[str, Any]] = {
             "examples": [
                 {