This is an automated email from the ASF dual-hosted git repository.
gopidesu pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new a480f77e747 Make task-sdk compatible with mypy 1.16.1 (#53047)
a480f77e747 is described below
commit a480f77e747a02a99e2afa50519e2ffe08952f1a
Author: GPK <[email protected]>
AuthorDate: Fri Jul 11 14:06:59 2025 +0100
Make task-sdk compatible with mypy 1.16.1 (#53047)
* Make task-sdk compatible with mypy 1.16.1
* Fixup mypy
* Fixup mypy
* Update comment for typeguard usage in secrets_masker
* Resolve comments
* Resolve pre-commit
---
task-sdk/src/airflow/sdk/definitions/_internal/expandinput.py | 5 ++++-
task-sdk/src/airflow/sdk/definitions/asset/__init__.py | 4 ++--
.../src/airflow/sdk/definitions/decorators/setup_teardown.py | 9 +++++++--
task-sdk/src/airflow/sdk/definitions/taskgroup.py | 8 +++++++-
task-sdk/src/airflow/sdk/execution_time/secrets_masker.py | 10 ++++------
task-sdk/src/airflow/sdk/execution_time/supervisor.py | 2 +-
6 files changed, 25 insertions(+), 13 deletions(-)
diff --git a/task-sdk/src/airflow/sdk/definitions/_internal/expandinput.py
b/task-sdk/src/airflow/sdk/definitions/_internal/expandinput.py
index 0424380e430..dd341bef1a1 100644
--- a/task-sdk/src/airflow/sdk/definitions/_internal/expandinput.py
+++ b/task-sdk/src/airflow/sdk/definitions/_internal/expandinput.py
@@ -203,11 +203,14 @@ class DictOfListsExpandInput(ResolveMixin):
# TODO: This initiates one API call for each XComArg. Would it be
# more efficient to do one single call and unpack the value here?
+
resolved = {
k: v.resolve(context) if _needs_run_time_resolution(v) else v for
k, v in self.value.items()
}
- all_lengths = self._get_map_lengths(resolved, upstream_map_indexes)
+ sized_resolved = {k: v for k, v in resolved.items() if isinstance(v,
Sized)}
+
+ all_lengths = self._get_map_lengths(sized_resolved,
upstream_map_indexes)
data = {k: self._expand_mapped_field(k, v, map_index, all_lengths) for
k, v in resolved.items()}
literal_keys = {k for k, _ in self._iter_parse_time_resolved_kwargs()}
diff --git a/task-sdk/src/airflow/sdk/definitions/asset/__init__.py
b/task-sdk/src/airflow/sdk/definitions/asset/__init__.py
index 90c972520a0..fae1ec79afe 100644
--- a/task-sdk/src/airflow/sdk/definitions/asset/__init__.py
+++ b/task-sdk/src/airflow/sdk/definitions/asset/__init__.py
@@ -633,7 +633,7 @@ class AssetBooleanCondition(BaseAsset):
class AssetAny(AssetBooleanCondition):
"""Use to combine assets schedule references in an "or" relationship."""
- agg_func = any
+ agg_func = any # type: ignore[assignment]
def __or__(self, other: BaseAsset) -> BaseAsset:
if not isinstance(other, BaseAsset):
@@ -656,7 +656,7 @@ class AssetAny(AssetBooleanCondition):
class AssetAll(AssetBooleanCondition):
"""Use to combine assets schedule references in an "and" relationship."""
- agg_func = all
+ agg_func = all # type: ignore[assignment]
def __and__(self, other: BaseAsset) -> BaseAsset:
if not isinstance(other, BaseAsset):
diff --git a/task-sdk/src/airflow/sdk/definitions/decorators/setup_teardown.py
b/task-sdk/src/airflow/sdk/definitions/decorators/setup_teardown.py
index bcb8e6d12e2..48cb7732494 100644
--- a/task-sdk/src/airflow/sdk/definitions/decorators/setup_teardown.py
+++ b/task-sdk/src/airflow/sdk/definitions/decorators/setup_teardown.py
@@ -18,7 +18,7 @@ from __future__ import annotations
import types
from collections.abc import Callable
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, cast
from airflow.exceptions import AirflowException
from airflow.sdk.bases.operator import BaseOperator
@@ -26,6 +26,7 @@ from airflow.sdk.definitions._internal.setup_teardown import
SetupTeardownContex
from airflow.sdk.definitions.decorators.task_group import _TaskGroupFactory
if TYPE_CHECKING:
+ from airflow.sdk.bases.decorator import _TaskDecorator
from airflow.sdk.definitions.xcom_arg import XComArg
try:
@@ -52,7 +53,8 @@ def setup_task(func: Callable) -> Callable:
func = python_task(func)
if isinstance(func, _TaskGroupFactory):
raise AirflowException("Task groups cannot be marked as setup or
teardown.")
- func.is_setup = True # type: ignore[attr-defined]
+ func = cast("_TaskDecorator", func)
+ func.is_setup = True # type: ignore[attr-defined] # TODO: Remove this
once mypy is bump to 1.16.1
return func
@@ -76,6 +78,9 @@ def teardown_task(_func=None, *, on_failure_fail_dagrun: bool
= False) -> Callab
func = python_task(func)
if isinstance(func, _TaskGroupFactory):
raise AirflowException("Task groups cannot be marked as setup or
teardown.")
+ func = cast("_TaskDecorator", func)
+
+ # TODO: Remove below attr-defined once mypy is bump to 1.16.1
func.is_teardown = True # type: ignore[attr-defined]
func.on_failure_fail_dagrun = on_failure_fail_dagrun # type:
ignore[attr-defined]
return func
diff --git a/task-sdk/src/airflow/sdk/definitions/taskgroup.py
b/task-sdk/src/airflow/sdk/definitions/taskgroup.py
index 6b2c2f65ee4..08db13de7a9 100644
--- a/task-sdk/src/airflow/sdk/definitions/taskgroup.py
+++ b/task-sdk/src/airflow/sdk/definitions/taskgroup.py
@@ -43,6 +43,7 @@ if TYPE_CHECKING:
from airflow.models.expandinput import SchedulerExpandInput
from airflow.sdk.bases.operator import BaseOperator
from airflow.sdk.definitions._internal.abstractoperator import
AbstractOperator
+ from airflow.sdk.definitions._internal.expandinput import
DictOfListsExpandInput, ListOfDictsExpandInput
from airflow.sdk.definitions._internal.mixins import DependencyMixin
from airflow.sdk.definitions.dag import DAG
from airflow.sdk.definitions.edges import EdgeModifier
@@ -613,7 +614,12 @@ class MappedTaskGroup(TaskGroup):
a ``@task_group`` function instead.
"""
- def __init__(self, *, expand_input: SchedulerExpandInput, **kwargs: Any)
-> None:
+ def __init__(
+ self,
+ *,
+ expand_input: SchedulerExpandInput | DictOfListsExpandInput |
ListOfDictsExpandInput,
+ **kwargs: Any,
+ ) -> None:
super().__init__(**kwargs)
self._expand_input = expand_input
diff --git a/task-sdk/src/airflow/sdk/execution_time/secrets_masker.py
b/task-sdk/src/airflow/sdk/execution_time/secrets_masker.py
index 28de9a2afa8..d800f16e615 100644
--- a/task-sdk/src/airflow/sdk/execution_time/secrets_masker.py
+++ b/task-sdk/src/airflow/sdk/execution_time/secrets_masker.py
@@ -27,13 +27,10 @@ from collections.abc import Callable, Generator, Iterable,
Iterator
from enum import Enum
from functools import cache, cached_property
from re import Pattern
-from typing import TYPE_CHECKING, Any, TextIO, TypeAlias, TypeVar
+from typing import Any, TextIO, TypeAlias, TypeVar
from airflow import settings
-if TYPE_CHECKING:
- from typing import TypeGuard
-
V1EnvVar = TypeVar("V1EnvVar")
Redactable: TypeAlias = str | V1EnvVar | dict[Any, Any] | tuple[Any, ...] |
list[Any]
Redacted: TypeAlias = Redactable | str
@@ -154,7 +151,8 @@ def _get_v1_env_var_type() -> type:
return V1EnvVar
-def _is_v1_env_var(v: Any) -> TypeGuard[V1EnvVar]:
+# TODO update return type to TypeGuard[V1EnvVar] once mypy 1.17.0 is available
+def _is_v1_env_var(v: Any) -> bool:
return isinstance(v, _get_v1_env_var_type())
@@ -256,7 +254,7 @@ class SecretsMasker(logging.Filter):
return to_return
if isinstance(item, Enum):
return self._redact(item=item.value, name=name, depth=depth,
max_depth=max_depth)
- if _is_v1_env_var(item):
+ if _is_v1_env_var(item) and hasattr(item, "to_dict"):
tmp: dict = item.to_dict() # type: ignore[attr-defined] #
V1EnvVar has a to_dict method
if should_hide_value_for_key(tmp.get("name", "")) and "value"
in tmp:
tmp["value"] = "***"
diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py
b/task-sdk/src/airflow/sdk/execution_time/supervisor.py
index f5eede6274d..5af102a786c 100644
--- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py
+++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py
@@ -666,7 +666,7 @@ class WatchedSubprocess:
return
# Escalation sequence: SIGINT -> SIGTERM -> SIGKILL
- escalation_path = [signal.SIGINT, signal.SIGTERM, signal.SIGKILL]
+ escalation_path: list[signal.Signals] = [signal.SIGINT,
signal.SIGTERM, signal.SIGKILL]
if force and signal_to_send in escalation_path:
# Start from `signal_to_send` and escalate to the end of the
escalation path
