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 bb3adb4b90 Improve importing the modules in Airflow cli package
(#33806)
bb3adb4b90 is described below
commit bb3adb4b90e3a601878bbec165015c2bcb2f73f0
Author: Hussein Awala <[email protected]>
AuthorDate: Mon Aug 28 00:10:16 2023 +0200
Improve importing the modules in Airflow cli package (#33806)
---
airflow/cli/cli_parser.py | 10 +++++++---
airflow/cli/commands/cheat_sheet_command.py | 7 +++++--
airflow/cli/commands/dag_command.py | 10 +++++++---
airflow/cli/commands/db_command.py | 6 +++++-
airflow/cli/commands/jobs_command.py | 6 +++++-
airflow/cli/commands/standalone_command.py | 5 ++++-
airflow/cli/commands/task_command.py | 10 +++++++---
airflow/cli/commands/webserver_command.py | 6 ++++--
airflow/cli/simple_table.py | 6 ++++--
airflow/cli/utils.py | 5 ++++-
tests/cli/commands/test_cheat_sheet_command.py | 2 +-
tests/cli/commands/test_task_command.py | 2 +-
tests/cli/test_cli_parser.py | 2 +-
13 files changed, 55 insertions(+), 22 deletions(-)
diff --git a/airflow/cli/cli_parser.py b/airflow/cli/cli_parser.py
index f3af605714..35051964b6 100644
--- a/airflow/cli/cli_parser.py
+++ b/airflow/cli/cli_parser.py
@@ -28,7 +28,7 @@ import collections
import logging
from argparse import Action
from functools import lru_cache
-from typing import Iterable
+from typing import TYPE_CHECKING, Iterable
import lazy_object_proxy
from rich_argparse import RawTextRichHelpFormatter, RichHelpFormatter
@@ -36,8 +36,6 @@ from rich_argparse import RawTextRichHelpFormatter,
RichHelpFormatter
from airflow.cli.cli_config import (
DAG_CLI_DICT,
ActionCommand,
- Arg,
- CLICommand,
DefaultHelpParser,
GroupCommand,
core_commands,
@@ -48,6 +46,12 @@ from airflow.executors.executor_loader import ExecutorLoader
from airflow.utils.helpers import partition
from airflow.www.extensions.init_auth_manager import get_auth_manager_cls
+if TYPE_CHECKING:
+ from airflow.cli.cli_config import (
+ Arg,
+ CLICommand,
+ )
+
airflow_commands = core_commands.copy() # make a copy to prevent bad
interactions in tests
log = logging.getLogger(__name__)
diff --git a/airflow/cli/commands/cheat_sheet_command.py
b/airflow/cli/commands/cheat_sheet_command.py
index 88d9c5940c..3e6789e813 100644
--- a/airflow/cli/commands/cheat_sheet_command.py
+++ b/airflow/cli/commands/cheat_sheet_command.py
@@ -16,12 +16,15 @@
# under the License.
from __future__ import annotations
-from typing import Iterable
+from typing import TYPE_CHECKING, Iterable
-from airflow.cli.cli_parser import ActionCommand, GroupCommand,
airflow_commands
+from airflow.cli.cli_parser import GroupCommand, airflow_commands
from airflow.cli.simple_table import AirflowConsole, SimpleTable
from airflow.utils.cli import suppress_logs_and_warning
+if TYPE_CHECKING:
+ from airflow.cli.cli_parser import ActionCommand
+
@suppress_logs_and_warning
def cheat_sheet(args):
diff --git a/airflow/cli/commands/dag_command.py
b/airflow/cli/commands/dag_command.py
index 4b54ac2f9f..8a5df9d028 100644
--- a/airflow/cli/commands/dag_command.py
+++ b/airflow/cli/commands/dag_command.py
@@ -26,10 +26,9 @@ import signal
import subprocess
import sys
import warnings
+from typing import TYPE_CHECKING
-from graphviz.dot import Dot
from sqlalchemy import delete, select
-from sqlalchemy.orm import Session
from airflow import settings
from airflow.api.client import get_current_api_client
@@ -41,7 +40,6 @@ from airflow.jobs.job import Job
from airflow.models import DagBag, DagModel, DagRun, TaskInstance
from airflow.models.dag import DAG
from airflow.models.serialized_dag import SerializedDagModel
-from airflow.timetables.base import DataInterval
from airflow.utils import cli as cli_utils, timezone
from airflow.utils.cli import get_dag, get_dags, process_subdir,
sigint_handler, suppress_logs_and_warning
from airflow.utils.dot_renderer import render_dag, render_dag_dependencies
@@ -49,6 +47,12 @@ from airflow.utils.providers_configuration_loader import
providers_configuration
from airflow.utils.session import NEW_SESSION, create_session, provide_session
from airflow.utils.state import DagRunState
+if TYPE_CHECKING:
+ from graphviz.dot import Dot
+ from sqlalchemy.orm import Session
+
+ from airflow.timetables.base import DataInterval
+
log = logging.getLogger(__name__)
diff --git a/airflow/cli/commands/db_command.py
b/airflow/cli/commands/db_command.py
index 449ba9c26c..1230202948 100644
--- a/airflow/cli/commands/db_command.py
+++ b/airflow/cli/commands/db_command.py
@@ -22,9 +22,10 @@ import os
import textwrap
import warnings
from tempfile import NamedTemporaryFile
+from typing import TYPE_CHECKING
from packaging.version import parse as parse_version
-from tenacity import RetryCallState, Retrying, stop_after_attempt, wait_fixed
+from tenacity import Retrying, stop_after_attempt, wait_fixed
from airflow import settings
from airflow.exceptions import AirflowException
@@ -34,6 +35,9 @@ from airflow.utils.db_cleanup import config_dict,
drop_archived_tables, export_a
from airflow.utils.process_utils import execute_interactive
from airflow.utils.providers_configuration_loader import
providers_configuration_loaded
+if TYPE_CHECKING:
+ from tenacity import RetryCallState
+
log = logging.getLogger(__name__)
diff --git a/airflow/cli/commands/jobs_command.py
b/airflow/cli/commands/jobs_command.py
index 3f22241db9..784e21fa79 100644
--- a/airflow/cli/commands/jobs_command.py
+++ b/airflow/cli/commands/jobs_command.py
@@ -16,8 +16,9 @@
# under the License.
from __future__ import annotations
+from typing import TYPE_CHECKING
+
from sqlalchemy import select
-from sqlalchemy.orm import Session
from airflow.jobs.job import Job
from airflow.utils.net import get_hostname
@@ -25,6 +26,9 @@ from airflow.utils.providers_configuration_loader import
providers_configuration
from airflow.utils.session import NEW_SESSION, provide_session
from airflow.utils.state import JobState
+if TYPE_CHECKING:
+ from sqlalchemy.orm import Session
+
@providers_configuration_loaded
@provide_session
diff --git a/airflow/cli/commands/standalone_command.py
b/airflow/cli/commands/standalone_command.py
index d502640743..3265b6b747 100644
--- a/airflow/cli/commands/standalone_command.py
+++ b/airflow/cli/commands/standalone_command.py
@@ -24,19 +24,22 @@ import subprocess
import threading
import time
from collections import deque
+from typing import TYPE_CHECKING
from termcolor import colored
from airflow.configuration import AIRFLOW_HOME, conf,
make_group_other_inaccessible
from airflow.executors import executor_constants
from airflow.executors.executor_loader import ExecutorLoader
-from airflow.jobs.base_job_runner import BaseJobRunner
from airflow.jobs.job import most_recent_job
from airflow.jobs.scheduler_job_runner import SchedulerJobRunner
from airflow.jobs.triggerer_job_runner import TriggererJobRunner
from airflow.utils import db
from airflow.utils.providers_configuration_loader import
providers_configuration_loaded
+if TYPE_CHECKING:
+ from airflow.jobs.base_job_runner import BaseJobRunner
+
class StandaloneCommand:
"""
diff --git a/airflow/cli/commands/task_command.py
b/airflow/cli/commands/task_command.py
index 8884f050fe..0cbf08e91d 100644
--- a/airflow/cli/commands/task_command.py
+++ b/airflow/cli/commands/task_command.py
@@ -25,13 +25,12 @@ import os
import sys
import textwrap
from contextlib import contextmanager, redirect_stderr, redirect_stdout,
suppress
-from typing import Generator, Protocol, Union, cast
+from typing import TYPE_CHECKING, Generator, Protocol, Union, cast
import pendulum
from pendulum.parsing.exceptions import ParserError
from sqlalchemy import select
from sqlalchemy.orm.exc import NoResultFound
-from sqlalchemy.orm.session import Session
from typing_extensions import Literal
from airflow import settings
@@ -45,7 +44,7 @@ from airflow.listeners.listener import get_listener_manager
from airflow.models import DagPickle, TaskInstance
from airflow.models.dag import DAG
from airflow.models.dagrun import DagRun
-from airflow.models.operator import Operator, needs_expansion
+from airflow.models.operator import needs_expansion
from airflow.models.param import ParamsDict
from airflow.models.taskinstance import TaskReturnCode
from airflow.settings import IS_K8S_EXECUTOR_POD
@@ -69,6 +68,11 @@ from airflow.utils.providers_configuration_loader import
providers_configuration
from airflow.utils.session import NEW_SESSION, create_session, provide_session
from airflow.utils.state import DagRunState
+if TYPE_CHECKING:
+ from sqlalchemy.orm.session import Session
+
+ from airflow.models.operator import Operator
+
log = logging.getLogger(__name__)
CreateIfNecessary = Union[Literal[False], Literal["db"], Literal["memory"]]
diff --git a/airflow/cli/commands/webserver_command.py
b/airflow/cli/commands/webserver_command.py
index 2aa686152c..d0b246eed0 100644
--- a/airflow/cli/commands/webserver_command.py
+++ b/airflow/cli/commands/webserver_command.py
@@ -24,10 +24,9 @@ import subprocess
import sys
import textwrap
import time
-import types
from contextlib import suppress
from time import sleep
-from typing import NoReturn
+from typing import TYPE_CHECKING, NoReturn
import daemon
import psutil
@@ -44,6 +43,9 @@ from airflow.utils.log.logging_mixin import LoggingMixin
from airflow.utils.process_utils import check_if_pidfile_process_is_running
from airflow.utils.providers_configuration_loader import
providers_configuration_loaded
+if TYPE_CHECKING:
+ import types
+
log = logging.getLogger(__name__)
diff --git a/airflow/cli/simple_table.py b/airflow/cli/simple_table.py
index f4f418742b..72750fbdc4 100644
--- a/airflow/cli/simple_table.py
+++ b/airflow/cli/simple_table.py
@@ -18,7 +18,7 @@ from __future__ import annotations
import inspect
import json
-from typing import Any, Callable, Sequence
+from typing import TYPE_CHECKING, Any, Callable, Sequence
from rich.box import ASCII_DOUBLE_HEAD
from rich.console import Console
@@ -27,10 +27,12 @@ from rich.table import Table
from tabulate import tabulate
from airflow.plugins_manager import PluginsDirectorySource
-from airflow.typing_compat import TypeGuard
from airflow.utils import yaml
from airflow.utils.platform import is_tty
+if TYPE_CHECKING:
+ from airflow.typing_compat import TypeGuard
+
def is_data_sequence(data: Sequence[dict | Any]) -> TypeGuard[Sequence[dict]]:
return all(isinstance(d, dict) for d in data)
diff --git a/airflow/cli/utils.py b/airflow/cli/utils.py
index a300798ec7..57a5348be2 100644
--- a/airflow/cli/utils.py
+++ b/airflow/cli/utils.py
@@ -17,8 +17,11 @@
from __future__ import annotations
-import io
import sys
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+ import io
class CliConflictError(Exception):
diff --git a/tests/cli/commands/test_cheat_sheet_command.py
b/tests/cli/commands/test_cheat_sheet_command.py
index c0fbcfe706..5d7b7f4190 100644
--- a/tests/cli/commands/test_cheat_sheet_command.py
+++ b/tests/cli/commands/test_cheat_sheet_command.py
@@ -21,7 +21,7 @@ import io
from unittest import mock
from airflow.cli import cli_parser
-from airflow.cli.cli_parser import ActionCommand, CLICommand, GroupCommand
+from airflow.cli.cli_config import ActionCommand, CLICommand, GroupCommand
def noop():
diff --git a/tests/cli/commands/test_task_command.py
b/tests/cli/commands/test_task_command.py
index e785b8ef46..c7937fa93f 100644
--- a/tests/cli/commands/test_task_command.py
+++ b/tests/cli/commands/test_task_command.py
@@ -501,7 +501,7 @@ class TestCliTasks:
assert 'echo "2022-01-08"' in output
@mock.patch("airflow.cli.commands.task_command.select")
- @mock.patch("airflow.cli.commands.task_command.Session.scalars")
+ @mock.patch("sqlalchemy.orm.session.Session.scalars")
@mock.patch("airflow.cli.commands.task_command.DagRun")
def test_task_render_with_custom_timetable(self, mock_dagrun,
mock_scalars, mock_select):
"""
diff --git a/tests/cli/test_cli_parser.py b/tests/cli/test_cli_parser.py
index 6235bc63a5..7eaf7febd3 100644
--- a/tests/cli/test_cli_parser.py
+++ b/tests/cli/test_cli_parser.py
@@ -158,7 +158,7 @@ class TestCli:
reload(cli_parser)
def test_falsy_default_value(self):
- arg = cli_parser.Arg(("--test",), default=0, type=int)
+ arg = cli_config.Arg(("--test",), default=0, type=int)
parser = argparse.ArgumentParser()
arg.add_to_parser(parser)
