fixup! fixup! style: fix provider mypy error

Lee-W · Lee-W · commit 39a20843e569 · 2026-04-06T19:10:02.000+08:00
diff --git a/airflow-core/src/airflow/assets/manager.py b/airflow-core/src/airflow/assets/manager.py
@@ -575,6 +575,8 @@ def resolve_asset_manager() -> AssetManager:
         key="asset_manager_kwargs",
         fallback={},
     )
+    if TYPE_CHECKING:
+        assert isinstance(_asset_manager_kwargs, dict)
     return _asset_manager_class(**_asset_manager_kwargs)
 
 
diff --git a/airflow-core/src/airflow/config_templates/airflow_local_settings.py b/airflow-core/src/airflow/config_templates/airflow_local_settings.py
@@ -20,7 +20,7 @@
 from __future__ import annotations
 
 import os
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, cast
 from urllib.parse import urlsplit
 
 from airflow.configuration import conf
@@ -159,59 +159,63 @@ def _default_conn_name_from(mod_path, hook_name):
             "logging/remote_task_handler_kwargs must be a JSON object (a python dict), we got "
             f"{type(remote_task_handler_kwargs)}"
         )
+    _handler_kwargs = cast("dict[str, Any]", remote_task_handler_kwargs)
     delete_local_copy = conf.getboolean("logging", "delete_local_logs")
 
     if remote_base_log_folder.startswith("s3://"):
         from airflow.providers.amazon.aws.log.s3_task_handler import S3RemoteLogIO
 
         _default_conn_name_from("airflow.providers.amazon.aws.hooks.s3", "S3Hook")
         REMOTE_TASK_LOG = S3RemoteLogIO(
-            **(
+            **cast(
+                "dict[str, Any]",
                 {
                     "base_log_folder": BASE_LOG_FOLDER,
                     "remote_base": remote_base_log_folder,
                     "delete_local_copy": delete_local_copy,
                 }
-                | remote_task_handler_kwargs
+                | _handler_kwargs,
             )
         )
-        remote_task_handler_kwargs = {}
+        _handler_kwargs = {}
 
     elif remote_base_log_folder.startswith("cloudwatch://"):
         from airflow.providers.amazon.aws.log.cloudwatch_task_handler import CloudWatchRemoteLogIO
 
         _default_conn_name_from("airflow.providers.amazon.aws.hooks.logs", "AwsLogsHook")
         url_parts = urlsplit(remote_base_log_folder)
         REMOTE_TASK_LOG = CloudWatchRemoteLogIO(
-            **(
+            **cast(
+                "dict[str, Any]",
                 {
                     "base_log_folder": BASE_LOG_FOLDER,
                     "remote_base": remote_base_log_folder,
                     "delete_local_copy": delete_local_copy,
                     "log_group_arn": url_parts.netloc + url_parts.path,
                 }
-                | remote_task_handler_kwargs
+                | _handler_kwargs,
             )
         )
-        remote_task_handler_kwargs = {}
+        _handler_kwargs = {}
     elif remote_base_log_folder.startswith("gs://"):
         from airflow.providers.google.cloud.log.gcs_task_handler import GCSRemoteLogIO
 
         _default_conn_name_from("airflow.providers.google.cloud.hooks.gcs", "GCSHook")
         key_path = conf.get_mandatory_value("logging", "google_key_path", fallback=None)
 
         REMOTE_TASK_LOG = GCSRemoteLogIO(
-            **(
+            **cast(
+                "dict[str, Any]",
                 {
                     "base_log_folder": BASE_LOG_FOLDER,
                     "remote_base": remote_base_log_folder,
                     "delete_local_copy": delete_local_copy,
                     "gcp_key_path": key_path,
                 }
-                | remote_task_handler_kwargs
+                | _handler_kwargs,
             )
         )
-        remote_task_handler_kwargs = {}
+        _handler_kwargs = {}
     elif remote_base_log_folder.startswith("wasb"):
         from airflow.providers.microsoft.azure.log.wasb_task_handler import WasbRemoteLogIO
 
@@ -224,17 +228,18 @@ def _default_conn_name_from(mod_path, hook_name):
         wasb_remote_base = remote_base_log_folder.removeprefix("wasb://")
 
         REMOTE_TASK_LOG = WasbRemoteLogIO(
-            **(
+            **cast(
+                "dict[str, Any]",
                 {
                     "base_log_folder": BASE_LOG_FOLDER,
                     "remote_base": wasb_remote_base,
                     "delete_local_copy": delete_local_copy,
                     "wasb_container": wasb_log_container,
                 }
-                | remote_task_handler_kwargs
+                | _handler_kwargs,
             )
         )
-        remote_task_handler_kwargs = {}
+        _handler_kwargs = {}
     elif remote_base_log_folder.startswith("stackdriver://"):
         key_path = conf.get_mandatory_value("logging", "GOOGLE_KEY_PATH", fallback=None)
         # stackdriver:///airflow-tasks => airflow-tasks
@@ -255,32 +260,34 @@ def _default_conn_name_from(mod_path, hook_name):
         _default_conn_name_from("airflow.providers.alibaba.cloud.hooks.oss", "OSSHook")
 
         REMOTE_TASK_LOG = OSSRemoteLogIO(
-            **(
+            **cast(
+                "dict[str, Any]",
                 {
                     "base_log_folder": BASE_LOG_FOLDER,
                     "remote_base": remote_base_log_folder,
                     "delete_local_copy": delete_local_copy,
                 }
-                | remote_task_handler_kwargs
+                | _handler_kwargs,
             )
         )
-        remote_task_handler_kwargs = {}
+        _handler_kwargs = {}
     elif remote_base_log_folder.startswith("hdfs://"):
         from airflow.providers.apache.hdfs.log.hdfs_task_handler import HdfsRemoteLogIO
 
         _default_conn_name_from("airflow.providers.apache.hdfs.hooks.webhdfs", "WebHDFSHook")
 
         REMOTE_TASK_LOG = HdfsRemoteLogIO(
-            **(
+            **cast(
+                "dict[str, Any]",
                 {
                     "base_log_folder": BASE_LOG_FOLDER,
                     "remote_base": urlsplit(remote_base_log_folder).path,
                     "delete_local_copy": delete_local_copy,
                 }
-                | remote_task_handler_kwargs
+                | _handler_kwargs,
             )
         )
-        remote_task_handler_kwargs = {}
+        _handler_kwargs = {}
     elif ELASTICSEARCH_HOST:
         from airflow.providers.elasticsearch.log.es_task_handler import ElasticsearchRemoteLogIO
 
diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_task_instances.py
@@ -21,7 +21,7 @@
 import itertools
 import os
 from datetime import timedelta
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 from unittest import mock
 
 import pendulum
@@ -151,7 +151,7 @@ def create_task_instances(
                 assert dag_version
 
             for mi in map_indexes:
-                kwargs = self.ti_init | {"map_index": mi}
+                kwargs: dict[str, Any] = self.ti_init | {"map_index": mi}
                 ti = TaskInstance(task=tasks[i], **kwargs, dag_version_id=dag_version.id)
                 session.add(ti)
                 ti.dag_run = dr
diff --git a/task-sdk/src/airflow/sdk/execution_time/task_runner.py b/task-sdk/src/airflow/sdk/execution_time/task_runner.py
@@ -225,28 +225,31 @@ def get_template_context(self) -> Context:
 
         # Cache the context object, which ensures that all calls to get_template_context
         # are operating on the same context object.
-        self._cached_template_context: Context = self._cached_template_context or {
-            # From the Task Execution interface
-            "dag": self.task.dag,
-            "inlets": self.task.inlets,
-            "map_index_template": self.task.map_index_template,
-            "outlets": self.task.outlets,
-            "run_id": self.run_id,
-            "task": self.task,
-            "task_instance": self,
-            "ti": self,
-            "outlet_events": OutletEventAccessors(),
-            "inlet_events": InletEventsAccessors(self.task.inlets),
-            "macros": MacrosAccessor(),
-            "params": validated_params,
-            # TODO: Make this go through Public API longer term.
-            # "test_mode": task_instance.test_mode,
-            "var": {
-                "json": VariableAccessor(deserialize_json=True),
-                "value": VariableAccessor(deserialize_json=False),
-            },
-            "conn": ConnectionAccessor(),
-        }
+        context = self._cached_template_context
+        if context is None:
+            context = {
+                # From the Task Execution interface
+                "dag": self.task.dag,
+                "inlets": self.task.inlets,
+                "map_index_template": self.task.map_index_template,
+                "outlets": self.task.outlets,
+                "run_id": self.run_id,
+                "task": self.task,
+                "task_instance": self,
+                "ti": self,
+                "outlet_events": OutletEventAccessors(),
+                "inlet_events": InletEventsAccessors(self.task.inlets),
+                "macros": MacrosAccessor(),
+                "params": validated_params,
+                # TODO: Make this go through Public API longer term.
+                # "test_mode": task_instance.test_mode,
+                "var": {
+                    "json": VariableAccessor(deserialize_json=True),
+                    "value": VariableAccessor(deserialize_json=False),
+                },
+                "conn": ConnectionAccessor(),
+            }
+            self._cached_template_context = context
         if from_server:
             dag_run = from_server.dag_run
             context_from_server: Context = {
@@ -265,7 +268,7 @@ def get_template_context(self) -> Context:
                     lambda: coerce_datetime(get_previous_dagrun_success(self.id).end_date)
                 ),
             }
-            self._cached_template_context.update(context_from_server)
+            context.update(context_from_server)
 
             if logical_date := coerce_datetime(dag_run.logical_date):
                 if TYPE_CHECKING:
@@ -276,7 +279,7 @@ def get_template_context(self) -> Context:
                 ts_nodash = logical_date.strftime("%Y%m%dT%H%M%S")
                 ts_nodash_with_tz = ts.replace("-", "").replace(":", "")
                 # logical_date and data_interval either coexist or be None together
-                self._cached_template_context.update(
+                context.update(
                     {
                         # keys that depend on logical_date
                         "logical_date": logical_date,
@@ -303,7 +306,7 @@ def get_template_context(self) -> Context:
             if upstream_map_indexes is not None:
                 setattr(self, "_upstream_map_indexes", upstream_map_indexes)
 
-        return self._cached_template_context
+        return context
 
     def render_templates(
         self, context: Context | None = None, jinja_env: jinja2.Environment | None = None
diff --git a/task-sdk/src/airflow/sdk/types.py b/task-sdk/src/airflow/sdk/types.py
@@ -21,7 +21,7 @@
 from collections.abc import Iterable
 from typing import TYPE_CHECKING, Any, NamedTuple, Protocol, TypeAlias
 
-from airflow.sdk.api.datamodels._generated import TerminalTIState, WeightRule
+from airflow.sdk.api.datamodels._generated import WeightRule
 from airflow.sdk.bases.xcom import BaseXCom
 from airflow.sdk.definitions._internal.types import NOTSET, ArgNotSet
 
@@ -160,7 +160,7 @@ def get_ti_count(
         task_group_id: str | None = None,
         logical_dates: list[AwareDatetime] | None = None,
         run_ids: list[str] | None = None,
-        states: list[str] | list[TerminalTIState] | None = None,
+        states: list[str] | None = None,
     ) -> int: ...
 
     @staticmethod

Original file line number	Diff line number	Diff line change
`@@ -575,6 +575,8 @@ def resolve_asset_manager() -> AssetManager:`
`575`	`575`	`key="asset_manager_kwargs",`
`576`	`576`	`fallback={},`
`577`	`577`	`)`
	`578`	`+ if TYPE_CHECKING:`
	`579`	`+ assert isinstance(_asset_manager_kwargs, dict)`
`578`	`580`	`return _asset_manager_class(**_asset_manager_kwargs)`
`579`	`581`
`580`	`582`
Original file line number	Diff line number	Diff line change
`@@ -20,7 +20,7 @@`
`20`	`20`	`from __future__ import annotations`
`21`	`21`
`22`	`22`	`import os`
`23`		`-from typing import TYPE_CHECKING, Any`
	`23`	`+from typing import TYPE_CHECKING, Any, cast`
`24`	`24`	`from urllib.parse import urlsplit`
`25`	`25`
`26`	`26`	`from airflow.configuration import conf`
`@@ -159,59 +159,63 @@ def _default_conn_name_from(mod_path, hook_name):`
`159`	`159`	`"logging/remote_task_handler_kwargs must be a JSON object (a python dict), we got "`
`160`	`160`	`f"{type(remote_task_handler_kwargs)}"`
`161`	`161`	`)`
	`162`	`+ _handler_kwargs = cast("dict[str, Any]", remote_task_handler_kwargs)`
`162`	`163`	`delete_local_copy = conf.getboolean("logging", "delete_local_logs")`
`163`	`164`
`164`	`165`	`if remote_base_log_folder.startswith("s3://"):`
`165`	`166`	`from airflow.providers.amazon.aws.log.s3_task_handler import S3RemoteLogIO`
`166`	`167`
`167`	`168`	`_default_conn_name_from("airflow.providers.amazon.aws.hooks.s3", "S3Hook")`
`168`	`169`	`REMOTE_TASK_LOG = S3RemoteLogIO(`
`169`		`- **(`
	`170`	`+ **cast(`
	`171`	`+ "dict[str, Any]",`
`170`	`172`	`{`
`171`	`173`	`"base_log_folder": BASE_LOG_FOLDER,`
`172`	`174`	`"remote_base": remote_base_log_folder,`
`173`	`175`	`"delete_local_copy": delete_local_copy,`
`174`	`176`	`}`
`175`		`- \| remote_task_handler_kwargs`
	`177`	`+ \| _handler_kwargs,`
`176`	`178`	`)`
`177`	`179`	`)`
`178`		`- remote_task_handler_kwargs = {}`
	`180`	`+ _handler_kwargs = {}`
`179`	`181`
`180`	`182`	`elif remote_base_log_folder.startswith("cloudwatch://"):`
`181`	`183`	`from airflow.providers.amazon.aws.log.cloudwatch_task_handler import CloudWatchRemoteLogIO`
`182`	`184`
`183`	`185`	`_default_conn_name_from("airflow.providers.amazon.aws.hooks.logs", "AwsLogsHook")`
`184`	`186`	`url_parts = urlsplit(remote_base_log_folder)`
`185`	`187`	`REMOTE_TASK_LOG = CloudWatchRemoteLogIO(`
`186`		`- **(`
	`188`	`+ **cast(`
	`189`	`+ "dict[str, Any]",`
`187`	`190`	`{`
`188`	`191`	`"base_log_folder": BASE_LOG_FOLDER,`
`189`	`192`	`"remote_base": remote_base_log_folder,`
`190`	`193`	`"delete_local_copy": delete_local_copy,`
`191`	`194`	`"log_group_arn": url_parts.netloc + url_parts.path,`
`192`	`195`	`}`
`193`		`- \| remote_task_handler_kwargs`
	`196`	`+ \| _handler_kwargs,`
`194`	`197`	`)`
`195`	`198`	`)`
`196`		`- remote_task_handler_kwargs = {}`
	`199`	`+ _handler_kwargs = {}`
`197`	`200`	`elif remote_base_log_folder.startswith("gs://"):`
`198`	`201`	`from airflow.providers.google.cloud.log.gcs_task_handler import GCSRemoteLogIO`
`199`	`202`
`200`	`203`	`_default_conn_name_from("airflow.providers.google.cloud.hooks.gcs", "GCSHook")`
`201`	`204`	`key_path = conf.get_mandatory_value("logging", "google_key_path", fallback=None)`
`202`	`205`
`203`	`206`	`REMOTE_TASK_LOG = GCSRemoteLogIO(`
`204`		`- **(`
	`207`	`+ **cast(`
	`208`	`+ "dict[str, Any]",`
`205`	`209`	`{`
`206`	`210`	`"base_log_folder": BASE_LOG_FOLDER,`
`207`	`211`	`"remote_base": remote_base_log_folder,`
`208`	`212`	`"delete_local_copy": delete_local_copy,`
`209`	`213`	`"gcp_key_path": key_path,`
`210`	`214`	`}`
`211`		`- \| remote_task_handler_kwargs`
	`215`	`+ \| _handler_kwargs,`
`212`	`216`	`)`
`213`	`217`	`)`
`214`		`- remote_task_handler_kwargs = {}`
	`218`	`+ _handler_kwargs = {}`
`215`	`219`	`elif remote_base_log_folder.startswith("wasb"):`
`216`	`220`	`from airflow.providers.microsoft.azure.log.wasb_task_handler import WasbRemoteLogIO`
`217`	`221`
`@@ -224,17 +228,18 @@ def _default_conn_name_from(mod_path, hook_name):`
`224`	`228`	`wasb_remote_base = remote_base_log_folder.removeprefix("wasb://")`
`225`	`229`
`226`	`230`	`REMOTE_TASK_LOG = WasbRemoteLogIO(`
`227`		`- **(`
	`231`	`+ **cast(`
	`232`	`+ "dict[str, Any]",`
`228`	`233`	`{`
`229`	`234`	`"base_log_folder": BASE_LOG_FOLDER,`
`230`	`235`	`"remote_base": wasb_remote_base,`
`231`	`236`	`"delete_local_copy": delete_local_copy,`
`232`	`237`	`"wasb_container": wasb_log_container,`
`233`	`238`	`}`
`234`		`- \| remote_task_handler_kwargs`
	`239`	`+ \| _handler_kwargs,`
`235`	`240`	`)`
`236`	`241`	`)`
`237`		`- remote_task_handler_kwargs = {}`
	`242`	`+ _handler_kwargs = {}`
`238`	`243`	`elif remote_base_log_folder.startswith("stackdriver://"):`
`239`	`244`	`key_path = conf.get_mandatory_value("logging", "GOOGLE_KEY_PATH", fallback=None)`
`240`	`245`	`# stackdriver:///airflow-tasks => airflow-tasks`
`@@ -255,32 +260,34 @@ def _default_conn_name_from(mod_path, hook_name):`
`255`	`260`	`_default_conn_name_from("airflow.providers.alibaba.cloud.hooks.oss", "OSSHook")`
`256`	`261`
`257`	`262`	`REMOTE_TASK_LOG = OSSRemoteLogIO(`
`258`		`- **(`
	`263`	`+ **cast(`
	`264`	`+ "dict[str, Any]",`
`259`	`265`	`{`
`260`	`266`	`"base_log_folder": BASE_LOG_FOLDER,`
`261`	`267`	`"remote_base": remote_base_log_folder,`
`262`	`268`	`"delete_local_copy": delete_local_copy,`
`263`	`269`	`}`
`264`		`- \| remote_task_handler_kwargs`
	`270`	`+ \| _handler_kwargs,`
`265`	`271`	`)`
`266`	`272`	`)`
`267`		`- remote_task_handler_kwargs = {}`
	`273`	`+ _handler_kwargs = {}`
`268`	`274`	`elif remote_base_log_folder.startswith("hdfs://"):`
`269`	`275`	`from airflow.providers.apache.hdfs.log.hdfs_task_handler import HdfsRemoteLogIO`
`270`	`276`
`271`	`277`	`_default_conn_name_from("airflow.providers.apache.hdfs.hooks.webhdfs", "WebHDFSHook")`
`272`	`278`
`273`	`279`	`REMOTE_TASK_LOG = HdfsRemoteLogIO(`
`274`		`- **(`
	`280`	`+ **cast(`
	`281`	`+ "dict[str, Any]",`
`275`	`282`	`{`
`276`	`283`	`"base_log_folder": BASE_LOG_FOLDER,`
`277`	`284`	`"remote_base": urlsplit(remote_base_log_folder).path,`
`278`	`285`	`"delete_local_copy": delete_local_copy,`
`279`	`286`	`}`
`280`		`- \| remote_task_handler_kwargs`
	`287`	`+ \| _handler_kwargs,`
`281`	`288`	`)`
`282`	`289`	`)`
`283`		`- remote_task_handler_kwargs = {}`
	`290`	`+ _handler_kwargs = {}`
`284`	`291`	`elif ELASTICSEARCH_HOST:`
`285`	`292`	`from airflow.providers.elasticsearch.log.es_task_handler import ElasticsearchRemoteLogIO`
`286`	`293`