From 333b4d30a809960322747fa11f7ca61392a42d2d Mon Sep 17 00:00:00 2001 From: PoAn Yang Date: Tue, 25 Aug 2026 09:58:52 +0900 Subject: [PATCH] Enable the airflow.utils.timezone ban rule Signed-off-by: PoAn Yang --- dev/airflow_perf/scheduler_dag_execution_timing.py | 2 +- devel-common/src/tests_common/pytest_plugin.py | 2 +- .../src/tests_common/test_utils/integration_setup.py | 6 +----- devel-common/src/tests_common/test_utils/version_compat.py | 2 +- .../tests/kubernetes_tests/test_kubernetes_pod_operator.py | 2 +- .../performance_dag/performance_dag_utils.py | 6 ++---- pyproject.toml | 3 +-- 7 files changed, 8 insertions(+), 15 deletions(-) diff --git a/dev/airflow_perf/scheduler_dag_execution_timing.py b/dev/airflow_perf/scheduler_dag_execution_timing.py index cfb4c4cecd654..2768ff28791f6 100755 --- a/dev/airflow_perf/scheduler_dag_execution_timing.py +++ b/dev/airflow_perf/scheduler_dag_execution_timing.py @@ -154,7 +154,7 @@ def create_dag_runs(dag, num_runs, session): """ Create `num_runs` of dag runs for sub-sequent schedules """ - from airflow.utils import timezone + from airflow.sdk import timezone from airflow.utils.state import DagRunState try: diff --git a/devel-common/src/tests_common/pytest_plugin.py b/devel-common/src/tests_common/pytest_plugin.py index b5e1df054e419..e2d9913205f11 100644 --- a/devel-common/src/tests_common/pytest_plugin.py +++ b/devel-common/src/tests_common/pytest_plugin.py @@ -3211,7 +3211,7 @@ def _import_timezone(): try: from airflow._shared.timezones import timezone except ImportError: - from airflow.utils import timezone + from airflow.utils import timezone # noqa: TID251 return timezone diff --git a/devel-common/src/tests_common/test_utils/integration_setup.py b/devel-common/src/tests_common/test_utils/integration_setup.py index 59476d8c4d06b..511c87223cedb 100644 --- a/devel-common/src/tests_common/test_utils/integration_setup.py +++ b/devel-common/src/tests_common/test_utils/integration_setup.py @@ -28,6 +28,7 @@ from airflow.dag_processing.dagbag import DagBag from airflow.models import DagRun from airflow.models.serialized_dag import SerializedDagModel +from airflow.sdk import timezone from airflow.serialization.definitions.dag import SerializedDAG from airflow.utils.session import create_session from airflow.utils.state import State @@ -35,11 +36,6 @@ from tests_common.test_utils.dag import create_scheduler_dag from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS, AIRFLOW_V_3_1_PLUS -try: - from airflow.sdk import timezone -except ImportError: - from airflow.utils import timezone # type: ignore[no-redef,attr-defined] - log = logging.getLogger(__name__) diff --git a/devel-common/src/tests_common/test_utils/version_compat.py b/devel-common/src/tests_common/test_utils/version_compat.py index e1b51ebe0347c..171dd06e791e6 100644 --- a/devel-common/src/tests_common/test_utils/version_compat.py +++ b/devel-common/src/tests_common/test_utils/version_compat.py @@ -53,7 +53,7 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]: XCOM_RETURN_KEY = BaseXCom.XCOM_RETURN_KEY else: from airflow.sensors.base import PokeReturnValue # type: ignore[no-redef] - from airflow.utils import timezone # type: ignore[attr-defined,no-redef] + from airflow.utils import timezone # type: ignore[attr-defined,no-redef] # noqa: TID251 from airflow.utils.decorators import remove_task_decorator # type: ignore[no-redef] from airflow.utils.types import NOTSET, ArgNotSet # type: ignore[attr-defined,no-redef] from airflow.utils.xcom import XCOM_RETURN_KEY # type: ignore[no-redef] diff --git a/kubernetes-tests/tests/kubernetes_tests/test_kubernetes_pod_operator.py b/kubernetes-tests/tests/kubernetes_tests/test_kubernetes_pod_operator.py index 97dc16925f0ad..bd14e9a4a612d 100644 --- a/kubernetes-tests/tests/kubernetes_tests/test_kubernetes_pod_operator.py +++ b/kubernetes-tests/tests/kubernetes_tests/test_kubernetes_pod_operator.py @@ -40,8 +40,8 @@ from airflow.providers.cncf.kubernetes.hooks.kubernetes import KubernetesHook from airflow.providers.cncf.kubernetes.operators.pod import KubernetesPodOperator from airflow.providers.cncf.kubernetes.utils.pod_manager import OnFinishAction, PodManager +from airflow.sdk import timezone from airflow.sdk.definitions.context import Context -from airflow.utils import timezone # type: ignore[attr-defined] from airflow.utils.types import DagRunType from airflow.version import version as airflow_version from kubernetes_tests.test_base import BaseK8STest, StringContainingId diff --git a/performance/src/performance_dags/performance_dag/performance_dag_utils.py b/performance/src/performance_dags/performance_dag/performance_dag_utils.py index 7dbf79c0e944d..e3de8a7383566 100644 --- a/performance/src/performance_dags/performance_dag/performance_dag_utils.py +++ b/performance/src/performance_dags/performance_dag/performance_dag_utils.py @@ -29,7 +29,7 @@ from datetime import datetime, timedelta from shutil import copyfile -import airflow +from airflow.sdk import timezone log = logging.getLogger(__name__) log.setLevel(logging.INFO) @@ -80,9 +80,7 @@ def add_perf_start_date_env_to_conf(performance_dag_conf: dict[str, str]) -> Non if "PERF_START_DATE" not in performance_dag_conf: start_ago = get_performance_dag_environment_variable(performance_dag_conf, "PERF_START_AGO") - perf_start_date = airflow.utils.timezone.utcnow - check_and_parse_time_delta( - "PERF_START_AGO", start_ago - ) + perf_start_date = timezone.utcnow() - check_and_parse_time_delta("PERF_START_AGO", start_ago) performance_dag_conf["PERF_START_DATE"] = str(perf_start_date) diff --git a/pyproject.toml b/pyproject.toml index 5179300f6f464..5b5b57192fba8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -910,8 +910,7 @@ banned-module-level-imports = ["numpy", "pandas", "polars"] # Note: we have to specify these all here (not in sub pyproject.toml files), else we'd have to duplicate these # all rules in each project https://github.com/astral-sh/ruff/issues/18723 -# Add this once providers are sorted -# "airflow.utils.timezone".msg = "Use airflow.sdk.timezone (or airflow._shared.timezones.timezone from within airflow-core)" +"airflow.utils.timezone".msg = "Use airflow.sdk.timezone (or airflow._shared.timezones.timezone from within airflow-core)" # "airflow_shared".msg = "Use airflow._shared or airflow.sdk._shared instead" # Deprecated in Python 3.11, Pending Removal in Python 3.15: https://github.com/python/cpython/issues/90817