diff --git a/providers/openlineage/tests/unit/openlineage/extractors/test_base.py b/providers/openlineage/tests/unit/openlineage/extractors/test_base.py index cbd852d393881..0aa1c996406a0 100644 --- a/providers/openlineage/tests/unit/openlineage/extractors/test_base.py +++ b/providers/openlineage/tests/unit/openlineage/extractors/test_base.py @@ -29,7 +29,7 @@ if AIRFLOW_V_3_0_PLUS: from airflow.sdk import BaseOperator else: - from airflow.models.baseoperator import BaseOperator + from airflow.models.baseoperator import BaseOperator # type: ignore[no-redef] from airflow.models.taskinstance import TaskInstanceState from airflow.providers.openlineage.extractors.base import ( diff --git a/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py b/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py index ae95f69df8d44..13aa9a06b9e61 100644 --- a/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py +++ b/providers/openlineage/tests/unit/openlineage/plugins/test_listener.py @@ -39,7 +39,7 @@ if AIRFLOW_V_3_0_PLUS: from airflow.sdk import BaseOperator else: - from airflow.models.baseoperator import BaseOperator + from airflow.models.baseoperator import BaseOperator # type: ignore[no-redef] from airflow.providers.openlineage.extractors.base import OperatorLineage from airflow.providers.openlineage.plugins.adapter import OpenLineageAdapter diff --git a/providers/openlineage/tests/unit/openlineage/utils/test_utils.py b/providers/openlineage/tests/unit/openlineage/utils/test_utils.py index 6357913e31665..839f192ed3a9b 100644 --- a/providers/openlineage/tests/unit/openlineage/utils/test_utils.py +++ b/providers/openlineage/tests/unit/openlineage/utils/test_utils.py @@ -26,16 +26,6 @@ from uuid6 import uuid7 from airflow import DAG -from airflow.utils import timezone - -from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS - -if AIRFLOW_V_3_0_PLUS: - from airflow.sdk import BaseOperator, task -else: - from airflow.decorators import task - from airflow.models.baseoperator import BaseOperator - from airflow.models.dagrun import DagRun from airflow.models.taskinstance import TaskInstance, TaskInstanceState from airflow.providers.common.compat.assets import Asset @@ -60,6 +50,7 @@ from airflow.serialization.serialized_objects import SerializedBaseOperator from airflow.timetables.events import EventsTimetable from airflow.timetables.trigger import CronTriggerTimetable +from airflow.utils import timezone from airflow.utils.state import DagRunState from airflow.utils.task_group import TaskGroup from airflow.utils.types import DagRunType @@ -68,6 +59,12 @@ from tests_common.test_utils.mock_operators import MockOperator from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS, AIRFLOW_V_3_1_PLUS +if AIRFLOW_V_3_0_PLUS: + from airflow.sdk import BaseOperator, task +else: + from airflow.decorators import task # type: ignore[no-redef] + from airflow.models.baseoperator import BaseOperator # type: ignore[no-redef] + BASH_OPERATOR_PATH = "airflow.providers.standard.operators.bash" PYTHON_OPERATOR_PATH = "airflow.providers.standard.operators.python" diff --git a/providers/standard/src/airflow/providers/standard/operators/bash.py b/providers/standard/src/airflow/providers/standard/operators/bash.py index afb1ad40a7eca..caee52072e10e 100644 --- a/providers/standard/src/airflow/providers/standard/operators/bash.py +++ b/providers/standard/src/airflow/providers/standard/operators/bash.py @@ -31,7 +31,7 @@ if AIRFLOW_V_3_1_PLUS: from airflow.sdk import BaseOperator else: - from airflow.models.baseoperator import BaseOperator + from airflow.models.baseoperator import BaseOperator # type: ignore[no-redef] if AIRFLOW_V_3_0_PLUS: from airflow.sdk.execution_time.context import context_to_airflow_vars diff --git a/providers/standard/src/airflow/providers/standard/operators/branch.py b/providers/standard/src/airflow/providers/standard/operators/branch.py index 44c958c4f2eea..064bb52afe188 100644 --- a/providers/standard/src/airflow/providers/standard/operators/branch.py +++ b/providers/standard/src/airflow/providers/standard/operators/branch.py @@ -27,12 +27,12 @@ if AIRFLOW_V_3_1_PLUS: from airflow.sdk import BaseOperator else: - from airflow.models.baseoperator import BaseOperator + from airflow.models.baseoperator import BaseOperator # type: ignore[no-redef] if AIRFLOW_V_3_0_PLUS: from airflow.providers.standard.utils.skipmixin import SkipMixin else: - from airflow.models.skipmixin import SkipMixin + from airflow.models.skipmixin import SkipMixin # type: ignore[no-redef] if TYPE_CHECKING: from airflow.sdk.definitions.context import Context diff --git a/providers/standard/src/airflow/providers/standard/operators/empty.py b/providers/standard/src/airflow/providers/standard/operators/empty.py index 4cbd6254d4a9e..cdfebf3a4139c 100644 --- a/providers/standard/src/airflow/providers/standard/operators/empty.py +++ b/providers/standard/src/airflow/providers/standard/operators/empty.py @@ -23,7 +23,7 @@ if AIRFLOW_V_3_1_PLUS: from airflow.sdk import BaseOperator else: - from airflow.models.baseoperator import BaseOperator + from airflow.models.baseoperator import BaseOperator # type: ignore[no-redef] if TYPE_CHECKING: from airflow.sdk.definitions.context import Context diff --git a/providers/standard/src/airflow/providers/standard/operators/python.py b/providers/standard/src/airflow/providers/standard/operators/python.py index 40bb51ba9d863..92387d28cbffd 100644 --- a/providers/standard/src/airflow/providers/standard/operators/python.py +++ b/providers/standard/src/airflow/providers/standard/operators/python.py @@ -60,13 +60,13 @@ if AIRFLOW_V_3_1_PLUS: from airflow.sdk import BaseOperator else: - from airflow.models.baseoperator import BaseOperator + from airflow.models.baseoperator import BaseOperator # type: ignore[no-redef] if AIRFLOW_V_3_0_PLUS: from airflow.providers.standard.operators.branch import BaseBranchOperator from airflow.providers.standard.utils.skipmixin import SkipMixin else: - from airflow.models.skipmixin import SkipMixin + from airflow.models.skipmixin import SkipMixin # type: ignore[no-redef] from airflow.operators.branch import BaseBranchOperator # type: ignore[no-redef] diff --git a/providers/standard/src/airflow/providers/standard/operators/smooth.py b/providers/standard/src/airflow/providers/standard/operators/smooth.py index 33f089421f741..6799e07a5ef05 100644 --- a/providers/standard/src/airflow/providers/standard/operators/smooth.py +++ b/providers/standard/src/airflow/providers/standard/operators/smooth.py @@ -24,7 +24,7 @@ if AIRFLOW_V_3_0_PLUS: from airflow.sdk import BaseOperator else: - from airflow.models.baseoperator import BaseOperator + from airflow.models.baseoperator import BaseOperator # type: ignore[no-redef] if TYPE_CHECKING: from airflow.sdk.definitions.context import Context diff --git a/providers/standard/src/airflow/providers/standard/sensors/external_task.py b/providers/standard/src/airflow/providers/standard/sensors/external_task.py index 12b41e7acf55b..974608bf714f4 100644 --- a/providers/standard/src/airflow/providers/standard/sensors/external_task.py +++ b/providers/standard/src/airflow/providers/standard/sensors/external_task.py @@ -58,8 +58,8 @@ from airflow.sdk import BaseOperator from airflow.sdk.definitions.context import Context else: - from airflow.models.baseoperator import BaseOperator - from airflow.utils.context import Context + from airflow.models.baseoperator import BaseOperator # type: ignore[no-redef] + from airflow.utils.context import Context # type: ignore[no-redef] if AIRFLOW_V_3_0_PLUS: diff --git a/providers/standard/src/airflow/providers/standard/utils/skipmixin.py b/providers/standard/src/airflow/providers/standard/utils/skipmixin.py index b7c8de2de2903..bd8e31bc72a34 100644 --- a/providers/standard/src/airflow/providers/standard/utils/skipmixin.py +++ b/providers/standard/src/airflow/providers/standard/utils/skipmixin.py @@ -49,8 +49,8 @@ def _ensure_tasks(nodes: Iterable[DAGNode]) -> Sequence[Operator]: from airflow.sdk import BaseOperator from airflow.sdk.definitions.mappedoperator import MappedOperator else: - from airflow.models.baseoperator import BaseOperator - from airflow.models.mappedoperator import MappedOperator + from airflow.models.baseoperator import BaseOperator # type: ignore[no-redef] + from airflow.models.mappedoperator import MappedOperator # type: ignore[no-redef] return [n for n in nodes if isinstance(n, BaseOperator | MappedOperator)] diff --git a/providers/standard/tests/unit/standard/operators/test_python.py b/providers/standard/tests/unit/standard/operators/test_python.py index 10e0bd986bf3b..e3ae3672718a2 100644 --- a/providers/standard/tests/unit/standard/operators/test_python.py +++ b/providers/standard/tests/unit/standard/operators/test_python.py @@ -47,13 +47,6 @@ AirflowProviderDeprecationWarning, DeserializingResultError, ) - -from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS - -if AIRFLOW_V_3_0_PLUS: - from airflow.sdk import BaseOperator -else: - from airflow.models.baseoperator import BaseOperator from airflow.models.taskinstance import TaskInstance, clear_task_instances, set_current_context from airflow.providers.standard.operators.empty import EmptyOperator from airflow.providers.standard.operators.python import ( @@ -78,6 +71,11 @@ from tests_common.test_utils.db import clear_db_runs from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_1, AIRFLOW_V_3_0_PLUS +if AIRFLOW_V_3_0_PLUS: + from airflow.sdk import BaseOperator +else: + from airflow.models.baseoperator import BaseOperator # type: ignore[no-redef] + if TYPE_CHECKING: from airflow.models.dag import DAG from airflow.models.dagrun import DagRun diff --git a/providers/standard/tests/unit/standard/sensors/test_external_task_sensor.py b/providers/standard/tests/unit/standard/sensors/test_external_task_sensor.py index 290f276091fb6..abbccf97cd9be 100644 --- a/providers/standard/tests/unit/standard/sensors/test_external_task_sensor.py +++ b/providers/standard/tests/unit/standard/sensors/test_external_task_sensor.py @@ -39,7 +39,7 @@ if AIRFLOW_V_3_0_PLUS: from airflow.sdk import BaseOperator else: - from airflow.models.baseoperator import BaseOperator + from airflow.models.baseoperator import BaseOperator # type: ignore[no-redef] from airflow.models.dag import DAG from airflow.models.serialized_dag import SerializedDagModel from airflow.models.xcom_arg import XComArg