diff --git a/airflow/utils/airflow_flask_app.py b/airflow/utils/airflow_flask_app.py index 751fcba8c360..f18f02a0ecb4 100644 --- a/airflow/utils/airflow_flask_app.py +++ b/airflow/utils/airflow_flask_app.py @@ -16,12 +16,13 @@ # under the License. from __future__ import annotations -from typing import Any, cast +from typing import TYPE_CHECKING, Any, cast from flask import Flask -from airflow.models.dagbag import DagBag -from airflow.www.extensions.init_appbuilder import AirflowAppBuilder +if TYPE_CHECKING: + from airflow.models.dagbag import DagBag + from airflow.www.extensions.init_appbuilder import AirflowAppBuilder class AirflowApp(Flask): diff --git a/airflow/utils/cli.py b/airflow/utils/cli.py index 0682e6531335..4343c797d924 100644 --- a/airflow/utils/cli.py +++ b/airflow/utils/cli.py @@ -33,7 +33,6 @@ import re2 from sqlalchemy import select -from sqlalchemy.orm import Session from airflow import settings from airflow.exceptions import AirflowException, RemovedInAirflow3Warning @@ -45,6 +44,8 @@ T = TypeVar("T", bound=Callable) if TYPE_CHECKING: + from sqlalchemy.orm import Session + from airflow.models.dag import DAG logger = logging.getLogger(__name__) diff --git a/airflow/utils/dag_edges.py b/airflow/utils/dag_edges.py index 7e31762ad91b..f7d0f7e7d832 100644 --- a/airflow/utils/dag_edges.py +++ b/airflow/utils/dag_edges.py @@ -16,9 +16,13 @@ # under the License. from __future__ import annotations -from airflow.models import Operator +from typing import TYPE_CHECKING + from airflow.models.abstractoperator import AbstractOperator -from airflow.models.dag import DAG + +if TYPE_CHECKING: + from airflow.models import Operator + from airflow.models.dag import DAG def dag_edges(dag: DAG): diff --git a/airflow/utils/db_cleanup.py b/airflow/utils/db_cleanup.py index f4e0293430d7..0f7eb8064ce0 100644 --- a/airflow/utils/db_cleanup.py +++ b/airflow/utils/db_cleanup.py @@ -27,24 +27,28 @@ import os from contextlib import contextmanager from dataclasses import dataclass -from typing import Any +from typing import TYPE_CHECKING, Any -from pendulum import DateTime from sqlalchemy import and_, column, false, func, inspect, select, table, text from sqlalchemy.exc import OperationalError, ProgrammingError from sqlalchemy.ext.compiler import compiles -from sqlalchemy.orm import Query, Session, aliased +from sqlalchemy.orm import aliased from sqlalchemy.sql.expression import ClauseElement, Executable, tuple_ from airflow import AirflowException from airflow.cli.simple_table import AirflowConsole from airflow.configuration import conf -from airflow.models import Base from airflow.utils import timezone from airflow.utils.db import reflect_tables from airflow.utils.helpers import ask_yesno from airflow.utils.session import NEW_SESSION, provide_session +if TYPE_CHECKING: + from pendulum import DateTime + from sqlalchemy.orm import Query, Session + + from airflow.models import Base + logger = logging.getLogger(__file__) ARCHIVE_TABLE_PREFIX = "_airflow_deleted__" diff --git a/airflow/utils/dot_renderer.py b/airflow/utils/dot_renderer.py index b6e2cdb7b094..05b91d9fc8ba 100644 --- a/airflow/utils/dot_renderer.py +++ b/airflow/utils/dot_renderer.py @@ -19,21 +19,23 @@ """Renderer DAG (tasks and dependencies) to the graphviz object.""" from __future__ import annotations -from typing import Any +from typing import TYPE_CHECKING, Any import graphviz from airflow import AirflowException -from airflow.models import TaskInstance from airflow.models.baseoperator import BaseOperator -from airflow.models.dag import DAG from airflow.models.mappedoperator import MappedOperator -from airflow.models.taskmixin import DependencyMixin -from airflow.serialization.serialized_objects import DagDependency from airflow.utils.dag_edges import dag_edges from airflow.utils.state import State from airflow.utils.task_group import TaskGroup +if TYPE_CHECKING: + from airflow.models import TaskInstance + from airflow.models.dag import DAG + from airflow.models.taskmixin import DependencyMixin + from airflow.serialization.serialized_objects import DagDependency + def _refine_color(color: str): """ diff --git a/airflow/utils/helpers.py b/airflow/utils/helpers.py index baef1d320902..acbe25bd92d2 100644 --- a/airflow/utils/helpers.py +++ b/airflow/utils/helpers.py @@ -30,7 +30,6 @@ from airflow.configuration import conf from airflow.exceptions import AirflowException, RemovedInAirflow3Warning -from airflow.utils.context import Context from airflow.utils.module_loading import import_string from airflow.utils.types import NOTSET @@ -38,6 +37,7 @@ import jinja2 from airflow.models.taskinstance import TaskInstance + from airflow.utils.context import Context KEY_REGEX = re.compile(r"^[\w.-]+$") GROUP_KEY_REGEX = re.compile(r"^[\w-]+$") diff --git a/airflow/utils/log/colored_log.py b/airflow/utils/log/colored_log.py index 3afa39410d04..2ed4fb4ed805 100644 --- a/airflow/utils/log/colored_log.py +++ b/airflow/utils/log/colored_log.py @@ -19,8 +19,7 @@ from __future__ import annotations import sys -from logging import LogRecord -from typing import Any +from typing import TYPE_CHECKING, Any import re2 from colorlog import TTYColoredFormatter @@ -28,6 +27,9 @@ from airflow.utils.log.timezone_aware import TimezoneAware +if TYPE_CHECKING: + from logging import LogRecord + DEFAULT_COLORS = { "DEBUG": "green", "INFO": "", diff --git a/airflow/utils/log/log_reader.py b/airflow/utils/log/log_reader.py index 9e751e47b1ab..589281fd55b4 100644 --- a/airflow/utils/log/log_reader.py +++ b/airflow/utils/log/log_reader.py @@ -19,17 +19,19 @@ import logging import time from functools import cached_property -from typing import Iterator - -from sqlalchemy.orm.session import Session +from typing import TYPE_CHECKING, Iterator from airflow.configuration import conf -from airflow.models.taskinstance import TaskInstance from airflow.utils.helpers import render_log_filename from airflow.utils.log.logging_mixin import ExternalLoggingMixin from airflow.utils.session import NEW_SESSION, provide_session from airflow.utils.state import TaskInstanceState +if TYPE_CHECKING: + from sqlalchemy.orm.session import Session + + from airflow.models.taskinstance import TaskInstance + class TaskLogReader: """Task log reader.""" diff --git a/airflow/utils/log/logging_mixin.py b/airflow/utils/log/logging_mixin.py index a3d3b4b00c71..5ef9885bdd13 100644 --- a/airflow/utils/log/logging_mixin.py +++ b/airflow/utils/log/logging_mixin.py @@ -22,11 +22,14 @@ import logging import sys from io import IOBase -from logging import Handler, Logger, StreamHandler -from typing import IO, Any, TypeVar, cast +from logging import Handler, StreamHandler +from typing import IO, TYPE_CHECKING, Any, TypeVar, cast import re2 +if TYPE_CHECKING: + from logging import Logger + # 7-bit C1 ANSI escape sequences ANSI_ESCAPE = re2.compile(r"\x1B[@-_][0-?]*[ -/]*[@-~]") diff --git a/airflow/utils/log/secrets_masker.py b/airflow/utils/log/secrets_masker.py index bb39dbf95b02..735d6e03a7a9 100644 --- a/airflow/utils/log/secrets_masker.py +++ b/airflow/utils/log/secrets_masker.py @@ -42,11 +42,12 @@ from airflow import settings from airflow.compat.functools import cache -from airflow.typing_compat import TypeGuard if TYPE_CHECKING: from kubernetes.client import V1EnvVar + from airflow.typing_compat import TypeGuard + Redactable = TypeVar("Redactable", str, "V1EnvVar", Dict[Any, Any], Tuple[Any, ...], List[Any]) Redacted = Union[Redactable, str] diff --git a/airflow/utils/log/trigger_handler.py b/airflow/utils/log/trigger_handler.py index 54f8ace65e74..42afcd8afeb5 100644 --- a/airflow/utils/log/trigger_handler.py +++ b/airflow/utils/log/trigger_handler.py @@ -22,8 +22,10 @@ from contextvars import ContextVar from copy import copy from logging.handlers import QueueHandler +from typing import TYPE_CHECKING -from airflow.utils.log.file_task_handler import FileTaskHandler +if TYPE_CHECKING: + from airflow.utils.log.file_task_handler import FileTaskHandler ctx_task_instance: ContextVar = ContextVar("task_instance") ctx_trigger_id: ContextVar = ContextVar("trigger_id") diff --git a/airflow/utils/mixins.py b/airflow/utils/mixins.py index e3c6a8efec6e..26e6c2192ca2 100644 --- a/airflow/utils/mixins.py +++ b/airflow/utils/mixins.py @@ -19,15 +19,14 @@ from __future__ import annotations import multiprocessing +import multiprocessing.context import typing from airflow.configuration import conf -from airflow.utils.context import Context if typing.TYPE_CHECKING: - import multiprocessing.context - from airflow.models.operator import Operator + from airflow.utils.context import Context class MultiprocessingStartMethodMixin: diff --git a/airflow/utils/module_loading.py b/airflow/utils/module_loading.py index d81ab65e8589..36d935f109ba 100644 --- a/airflow/utils/module_loading.py +++ b/airflow/utils/module_loading.py @@ -19,8 +19,10 @@ import pkgutil from importlib import import_module -from types import ModuleType -from typing import Callable +from typing import TYPE_CHECKING, Callable + +if TYPE_CHECKING: + from types import ModuleType def import_string(dotted_path: str): diff --git a/airflow/utils/session.py b/airflow/utils/session.py index 400cac7375f9..b3b610d19952 100644 --- a/airflow/utils/session.py +++ b/airflow/utils/session.py @@ -21,12 +21,14 @@ from inspect import signature from typing import Callable, Generator, TypeVar, cast +from sqlalchemy.orm import Session as SASession + from airflow import settings from airflow.typing_compat import ParamSpec @contextlib.contextmanager -def create_session() -> Generator[settings.SASession, None, None]: +def create_session() -> Generator[SASession, None, None]: """Contextmanager that will create and teardown a session.""" Session = getattr(settings, "Session", None) if Session is None: @@ -83,4 +85,4 @@ def wrapper(*args, **kwargs) -> RT: # the 'session' argument to be of type Session instead of Session | None, # making it easier to type hint the function body without dealing with the None # case that can never happen at runtime. -NEW_SESSION: settings.SASession = cast(settings.SASession, None) +NEW_SESSION: SASession = cast(SASession, None) diff --git a/airflow/utils/sqlalchemy.py b/airflow/utils/sqlalchemy.py index 1a1331427f4d..9166a11244cc 100644 --- a/airflow/utils/sqlalchemy.py +++ b/airflow/utils/sqlalchemy.py @@ -28,9 +28,8 @@ from dateutil import relativedelta from sqlalchemy import TIMESTAMP, PickleType, and_, event, false, nullsfirst, or_, true, tuple_ from sqlalchemy.dialects import mssql, mysql -from sqlalchemy.exc import OperationalError -from sqlalchemy.sql import ColumnElement, Select -from sqlalchemy.types import JSON, Text, TypeDecorator, TypeEngine, UnicodeText +from sqlalchemy.sql import Select +from sqlalchemy.types import JSON, Text, TypeDecorator, UnicodeText from airflow import settings from airflow.configuration import conf @@ -39,8 +38,11 @@ if TYPE_CHECKING: from kubernetes.client.models.v1_pod import V1Pod + from sqlalchemy.exc import OperationalError from sqlalchemy.orm import Query, Session + from sqlalchemy.sql import ColumnElement from sqlalchemy.sql.expression import ColumnOperators + from sqlalchemy.types import TypeEngine log = logging.getLogger(__name__) diff --git a/airflow/utils/task_group.py b/airflow/utils/task_group.py index 841abc710678..e892f0e85dbb 100644 --- a/airflow/utils/task_group.py +++ b/airflow/utils/task_group.py @@ -33,7 +33,7 @@ DuplicateTaskIdFound, TaskAlreadyInTaskGroup, ) -from airflow.models.taskmixin import DAGNode, DependencyMixin +from airflow.models.taskmixin import DAGNode from airflow.serialization.enums import DagAttributeTypes from airflow.utils.helpers import validate_group_key @@ -45,6 +45,7 @@ from airflow.models.dag import DAG from airflow.models.expandinput import ExpandInput from airflow.models.operator import Operator + from airflow.models.taskmixin import DependencyMixin from airflow.utils.edgemodifier import EdgeModifier