This is an automated email from the ASF dual-hosted git repository.
husseinawala 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 58d8577f35 Improve importing the module in Airflow utils package
(#33803)
58d8577f35 is described below
commit 58d8577f3528e7e143addefb03711bdc82cb1782
Author: Hussein Awala <[email protected]>
AuthorDate: Sun Aug 27 22:00:10 2023 +0200
Improve importing the module in Airflow utils package (#33803)
---
airflow/utils/airflow_flask_app.py | 7 ++++---
airflow/utils/cli.py | 3 ++-
airflow/utils/dag_edges.py | 8 ++++++--
airflow/utils/db_cleanup.py | 12 ++++++++----
airflow/utils/dot_renderer.py | 12 +++++++-----
airflow/utils/helpers.py | 2 +-
airflow/utils/log/colored_log.py | 6 ++++--
airflow/utils/log/log_reader.py | 10 ++++++----
airflow/utils/log/logging_mixin.py | 7 +++++--
airflow/utils/log/secrets_masker.py | 3 ++-
airflow/utils/log/trigger_handler.py | 4 +++-
airflow/utils/mixins.py | 5 ++---
airflow/utils/module_loading.py | 6 ++++--
airflow/utils/session.py | 6 ++++--
airflow/utils/sqlalchemy.py | 8 +++++---
airflow/utils/task_group.py | 3 ++-
16 files changed, 65 insertions(+), 37 deletions(-)
diff --git a/airflow/utils/airflow_flask_app.py
b/airflow/utils/airflow_flask_app.py
index 751fcba8c3..f18f02a0ec 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 0682e65313..4343c797d9 100644
--- a/airflow/utils/cli.py
+++ b/airflow/utils/cli.py
@@ -33,7 +33,6 @@ from typing import TYPE_CHECKING, Callable, TypeVar, cast
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 @@ from airflow.utils.session import NEW_SESSION, provide_session
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 7e31762ad9..f7d0f7e7d8 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 f4e0293430..0f7eb8064c 100644
--- a/airflow/utils/db_cleanup.py
+++ b/airflow/utils/db_cleanup.py
@@ -27,24 +27,28 @@ import logging
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 b6e2cdb7b0..05b91d9fc8 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 baef1d3209..acbe25bd92 100644
--- a/airflow/utils/helpers.py
+++ b/airflow/utils/helpers.py
@@ -30,7 +30,6 @@ from lazy_object_proxy import Proxy
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 @@ if TYPE_CHECKING:
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 3afa39410d..2ed4fb4ed8 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 colorlog.escape_codes import esc, escape_codes
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 9e751e47b1..589281fd55 100644
--- a/airflow/utils/log/log_reader.py
+++ b/airflow/utils/log/log_reader.py
@@ -19,17 +19,19 @@ from __future__ import annotations
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 a3d3b4b00c..5ef9885bdd 100644
--- a/airflow/utils/log/logging_mixin.py
+++ b/airflow/utils/log/logging_mixin.py
@@ -22,11 +22,14 @@ import enum
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 bb39dbf95b..735d6e03a7 100644
--- a/airflow/utils/log/secrets_masker.py
+++ b/airflow/utils/log/secrets_masker.py
@@ -42,11 +42,12 @@ import re2
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 54f8ace65e..42afcd8afe 100644
--- a/airflow/utils/log/trigger_handler.py
+++ b/airflow/utils/log/trigger_handler.py
@@ -22,8 +22,10 @@ import logging
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 e3c6a8efec..26e6c2192c 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 d81ab65e85..36d935f109 100644
--- a/airflow/utils/module_loading.py
+++ b/airflow/utils/module_loading.py
@@ -19,8 +19,10 @@ from __future__ import annotations
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 400cac7375..b3b610d199 100644
--- a/airflow/utils/session.py
+++ b/airflow/utils/session.py
@@ -21,12 +21,14 @@ from functools import wraps
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 provide_session(func: Callable[PS, RT]) -> Callable[PS,
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 1a1331427f..9166a11244 100644
--- a/airflow/utils/sqlalchemy.py
+++ b/airflow/utils/sqlalchemy.py
@@ -28,9 +28,8 @@ import pendulum
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 @@ from airflow.utils.timezone import make_naive
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 841abc7106..e892f0e85d 100644
--- a/airflow/utils/task_group.py
+++ b/airflow/utils/task_group.py
@@ -33,7 +33,7 @@ from airflow.exceptions import (
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 @@ if TYPE_CHECKING:
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
