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 a23f3181fe Improve importing the module in Airflow decorators package
(#33804)
a23f3181fe is described below
commit a23f3181fec999194dbc6b2215701a98eb34700c
Author: Hussein Awala <[email protected]>
AuthorDate: Sun Aug 27 22:12:29 2023 +0200
Improve importing the module in Airflow decorators package (#33804)
---
airflow/decorators/base.py | 26 ++++++++++++++++++--------
airflow/decorators/branch_python.py | 7 +++++--
airflow/decorators/external_python.py | 7 +++++--
airflow/decorators/python.py | 7 +++++--
airflow/decorators/python_virtualenv.py | 7 +++++--
airflow/decorators/sensor.py | 7 +++++--
airflow/decorators/setup_teardown.py | 7 +++++--
airflow/decorators/short_circuit.py | 7 +++++--
airflow/decorators/task_group.py | 6 ++++--
9 files changed, 57 insertions(+), 24 deletions(-)
diff --git a/airflow/decorators/base.py b/airflow/decorators/base.py
index 750e1fa1e7..d601483a3a 100644
--- a/airflow/decorators/base.py
+++ b/airflow/decorators/base.py
@@ -22,6 +22,7 @@ import warnings
from functools import cached_property
from textwrap import dedent
from typing import (
+ TYPE_CHECKING,
Any,
Callable,
ClassVar,
@@ -39,7 +40,6 @@ from typing import (
import attr
import re2
import typing_extensions
-from sqlalchemy.orm import Session
from airflow import Dataset
from airflow.exceptions import AirflowException
@@ -51,28 +51,38 @@ from airflow.models.baseoperator import (
get_merged_defaults,
parse_retries,
)
-from airflow.models.dag import DAG, DagContext
+from airflow.models.dag import DagContext
from airflow.models.expandinput import (
EXPAND_INPUT_EMPTY,
DictOfListsExpandInput,
- ExpandInput,
ListOfDictsExpandInput,
- OperatorExpandArgument,
- OperatorExpandKwargsArgument,
is_mappable,
)
-from airflow.models.mappedoperator import MappedOperator, ValidationSource,
ensure_xcomarg_return_value
+from airflow.models.mappedoperator import MappedOperator,
ensure_xcomarg_return_value
from airflow.models.pool import Pool
from airflow.models.xcom_arg import XComArg
from airflow.typing_compat import ParamSpec, Protocol
from airflow.utils import timezone
-from airflow.utils.context import KNOWN_CONTEXT_KEYS, Context
+from airflow.utils.context import KNOWN_CONTEXT_KEYS
from airflow.utils.decorators import remove_task_decorator
from airflow.utils.helpers import prevent_duplicates
-from airflow.utils.task_group import TaskGroup, TaskGroupContext
+from airflow.utils.task_group import TaskGroupContext
from airflow.utils.trigger_rule import TriggerRule
from airflow.utils.types import NOTSET
+if TYPE_CHECKING:
+ from sqlalchemy.orm import Session
+
+ from airflow.models.dag import DAG
+ from airflow.models.expandinput import (
+ ExpandInput,
+ OperatorExpandArgument,
+ OperatorExpandKwargsArgument,
+ )
+ from airflow.models.mappedoperator import ValidationSource
+ from airflow.utils.context import Context
+ from airflow.utils.task_group import TaskGroup
+
class ExpandableFactory(Protocol):
"""Protocol providing inspection against wrapped function.
diff --git a/airflow/decorators/branch_python.py
b/airflow/decorators/branch_python.py
index 39105b0072..3ac11f0efa 100644
--- a/airflow/decorators/branch_python.py
+++ b/airflow/decorators/branch_python.py
@@ -16,12 +16,15 @@
# under the License.
from __future__ import annotations
-from typing import Callable
+from typing import TYPE_CHECKING, Callable
-from airflow.decorators.base import TaskDecorator, task_decorator_factory
+from airflow.decorators.base import task_decorator_factory
from airflow.decorators.python import _PythonDecoratedOperator
from airflow.operators.python import BranchPythonOperator
+if TYPE_CHECKING:
+ from airflow.decorators.base import TaskDecorator
+
class _BranchPythonDecoratedOperator(_PythonDecoratedOperator,
BranchPythonOperator):
"""Wraps a Python callable and captures args/kwargs when called for
execution."""
diff --git a/airflow/decorators/external_python.py
b/airflow/decorators/external_python.py
index f17a9f3f1c..1e39ed561b 100644
--- a/airflow/decorators/external_python.py
+++ b/airflow/decorators/external_python.py
@@ -16,12 +16,15 @@
# under the License.
from __future__ import annotations
-from typing import Callable
+from typing import TYPE_CHECKING, Callable
-from airflow.decorators.base import TaskDecorator, task_decorator_factory
+from airflow.decorators.base import task_decorator_factory
from airflow.decorators.python import _PythonDecoratedOperator
from airflow.operators.python import ExternalPythonOperator
+if TYPE_CHECKING:
+ from airflow.decorators.base import TaskDecorator
+
class _PythonExternalDecoratedOperator(_PythonDecoratedOperator,
ExternalPythonOperator):
"""Wraps a Python callable and captures args/kwargs when called for
execution."""
diff --git a/airflow/decorators/python.py b/airflow/decorators/python.py
index d4423a2092..7a890cf862 100644
--- a/airflow/decorators/python.py
+++ b/airflow/decorators/python.py
@@ -16,11 +16,14 @@
# under the License.
from __future__ import annotations
-from typing import Callable, Sequence
+from typing import TYPE_CHECKING, Callable, Sequence
-from airflow.decorators.base import DecoratedOperator, TaskDecorator,
task_decorator_factory
+from airflow.decorators.base import DecoratedOperator, task_decorator_factory
from airflow.operators.python import PythonOperator
+if TYPE_CHECKING:
+ from airflow.decorators.base import TaskDecorator
+
class _PythonDecoratedOperator(DecoratedOperator, PythonOperator):
"""
diff --git a/airflow/decorators/python_virtualenv.py
b/airflow/decorators/python_virtualenv.py
index 9b632e0c87..2eb8678779 100644
--- a/airflow/decorators/python_virtualenv.py
+++ b/airflow/decorators/python_virtualenv.py
@@ -16,12 +16,15 @@
# under the License.
from __future__ import annotations
-from typing import Callable
+from typing import TYPE_CHECKING, Callable
-from airflow.decorators.base import TaskDecorator, task_decorator_factory
+from airflow.decorators.base import task_decorator_factory
from airflow.decorators.python import _PythonDecoratedOperator
from airflow.operators.python import PythonVirtualenvOperator
+if TYPE_CHECKING:
+ from airflow.decorators.base import TaskDecorator
+
class _PythonVirtualenvDecoratedOperator(_PythonDecoratedOperator,
PythonVirtualenvOperator):
"""Wraps a Python callable and captures args/kwargs when called for
execution."""
diff --git a/airflow/decorators/sensor.py b/airflow/decorators/sensor.py
index e8363317db..1d6ab86c9f 100644
--- a/airflow/decorators/sensor.py
+++ b/airflow/decorators/sensor.py
@@ -17,11 +17,14 @@
from __future__ import annotations
-from typing import Callable, Sequence
+from typing import TYPE_CHECKING, Callable, Sequence
-from airflow.decorators.base import TaskDecorator, get_unique_task_id,
task_decorator_factory
+from airflow.decorators.base import get_unique_task_id, task_decorator_factory
from airflow.sensors.python import PythonSensor
+if TYPE_CHECKING:
+ from airflow.decorators.base import TaskDecorator
+
class DecoratedSensorOperator(PythonSensor):
"""
diff --git a/airflow/decorators/setup_teardown.py
b/airflow/decorators/setup_teardown.py
index e9bb19b5c3..ecb3dc289f 100644
--- a/airflow/decorators/setup_teardown.py
+++ b/airflow/decorators/setup_teardown.py
@@ -17,14 +17,17 @@
from __future__ import annotations
import types
-from typing import Callable
+from typing import TYPE_CHECKING, Callable
-from airflow import AirflowException, XComArg
+from airflow import AirflowException
from airflow.decorators import python_task
from airflow.decorators.task_group import _TaskGroupFactory
from airflow.models import BaseOperator
from airflow.utils.setup_teardown import SetupTeardownContext
+if TYPE_CHECKING:
+ from airflow import XComArg
+
def setup_task(func: Callable) -> Callable:
# Using FunctionType here since _TaskDecorator is also a callable
diff --git a/airflow/decorators/short_circuit.py
b/airflow/decorators/short_circuit.py
index dd94daddd6..210a0e0453 100644
--- a/airflow/decorators/short_circuit.py
+++ b/airflow/decorators/short_circuit.py
@@ -16,12 +16,15 @@
# under the License.
from __future__ import annotations
-from typing import Callable
+from typing import TYPE_CHECKING, Callable
-from airflow.decorators.base import TaskDecorator, task_decorator_factory
+from airflow.decorators.base import task_decorator_factory
from airflow.decorators.python import _PythonDecoratedOperator
from airflow.operators.python import ShortCircuitOperator
+if TYPE_CHECKING:
+ from airflow.decorators.base import TaskDecorator
+
class _ShortCircuitDecoratedOperator(_PythonDecoratedOperator,
ShortCircuitOperator):
"""Wraps a Python callable and captures args/kwargs when called for
execution."""
diff --git a/airflow/decorators/task_group.py b/airflow/decorators/task_group.py
index f0d510193c..312c174b0e 100644
--- a/airflow/decorators/task_group.py
+++ b/airflow/decorators/task_group.py
@@ -36,8 +36,6 @@ from airflow.models.expandinput import (
DictOfListsExpandInput,
ListOfDictsExpandInput,
MappedArgument,
- OperatorExpandArgument,
- OperatorExpandKwargsArgument,
)
from airflow.models.taskmixin import DAGNode
from airflow.models.xcom_arg import XComArg
@@ -47,6 +45,10 @@ from airflow.utils.task_group import MappedTaskGroup,
TaskGroup
if TYPE_CHECKING:
from airflow.models.dag import DAG
+ from airflow.models.expandinput import (
+ OperatorExpandArgument,
+ OperatorExpandKwargsArgument,
+ )
FParams = ParamSpec("FParams")
FReturn = TypeVar("FReturn", None, DAGNode)
