From cf752bc8604051df52d80512588441decc10aaf4 Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Wed, 12 Nov 2025 16:09:27 +0800 Subject: [PATCH] Improve is_container annotation Using TypeGuard lets the function perform type-narrowing for us, so usages after it would not need a custom cast. The downside is we now need to specify a bound, but we don't actually use this function that often anyway and can just list all the possibilities. --- airflow-core/src/airflow/utils/helpers.py | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) diff --git a/airflow-core/src/airflow/utils/helpers.py b/airflow-core/src/airflow/utils/helpers.py index 112b6e70c9866..4d79a28d41e77 100644 --- a/airflow-core/src/airflow/utils/helpers.py +++ b/airflow-core/src/airflow/utils/helpers.py @@ -23,7 +23,7 @@ import signal from collections.abc import Callable, Generator, Iterable, MutableMapping from functools import cache -from typing import TYPE_CHECKING, Any, TypeVar, cast +from typing import TYPE_CHECKING, Any, TypeVar, cast, overload from urllib.parse import urljoin from lazy_object_proxy import Proxy @@ -33,11 +33,16 @@ from airflow.utils.types import NOTSET if TYPE_CHECKING: + from datetime import datetime + from typing import TypeGuard + import jinja2 from airflow.models.taskinstance import TaskInstance from airflow.sdk.definitions.context import Context + CT = TypeVar("CT", str, datetime) + KEY_REGEX = re.compile(r"^[\w.-]+$") GROUP_KEY_REGEX = re.compile(r"^[\w-]+$") CAMELCASE_TO_SNAKE_CASE_REGEX = re.compile(r"(?!^)([A-Z]+)") @@ -90,7 +95,15 @@ def handler(signum, frame): signal.alarm(0) -def is_container(obj: Any) -> bool: +@overload +def is_container(obj: None | int | Iterable[int] | range) -> TypeGuard[Iterable[int]]: ... + + +@overload +def is_container(obj: None | CT | Iterable[CT]) -> TypeGuard[Iterable[CT]]: ... + + +def is_container(obj) -> bool: """Test if an object is a container (iterable) but not a string.""" if isinstance(obj, Proxy): # Proxy of any object is considered a container because it implements __iter__