diff --git a/providers/ftp/src/airflow/providers/ftp/hooks/ftp.py b/providers/ftp/src/airflow/providers/ftp/hooks/ftp.py index 3fd389c57e6ba..010d33e8a15f7 100644 --- a/providers/ftp/src/airflow/providers/ftp/hooks/ftp.py +++ b/providers/ftp/src/airflow/providers/ftp/hooks/ftp.py @@ -23,10 +23,7 @@ from collections.abc import Callable from typing import Any, cast -try: - from airflow.sdk import BaseHook -except ImportError: - from airflow.hooks.base import BaseHook # type: ignore[attr-defined,no-redef] +from airflow.providers.ftp.version_compat import BaseHook logger = logging.getLogger(__name__) diff --git a/providers/ftp/src/airflow/providers/ftp/sensors/ftp.py b/providers/ftp/src/airflow/providers/ftp/sensors/ftp.py index c2b69fadaa36b..f0e3f31d5ba8c 100644 --- a/providers/ftp/src/airflow/providers/ftp/sensors/ftp.py +++ b/providers/ftp/src/airflow/providers/ftp/sensors/ftp.py @@ -23,12 +23,7 @@ from typing import TYPE_CHECKING from airflow.providers.ftp.hooks.ftp import FTPHook, FTPSHook -from airflow.providers.ftp.version_compat import AIRFLOW_V_3_0_PLUS - -if AIRFLOW_V_3_0_PLUS: - from airflow.sdk import BaseSensorOperator -else: - from airflow.sensors.base import BaseSensorOperator # type: ignore[no-redef] +from airflow.providers.ftp.version_compat import BaseSensorOperator if TYPE_CHECKING: try: diff --git a/providers/ftp/src/airflow/providers/ftp/version_compat.py b/providers/ftp/src/airflow/providers/ftp/version_compat.py index 4f8d5e32bca4a..680f6501ec916 100644 --- a/providers/ftp/src/airflow/providers/ftp/version_compat.py +++ b/providers/ftp/src/airflow/providers/ftp/version_compat.py @@ -33,10 +33,17 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]: AIRFLOW_V_3_0_PLUS = get_base_airflow_version_tuple() >= (3, 0, 0) +AIRFLOW_V_3_1_PLUS: bool = get_base_airflow_version_tuple() >= (3, 1, 0) if AIRFLOW_V_3_0_PLUS: - from airflow.sdk import BaseOperator + from airflow.sdk import BaseOperator, BaseSensorOperator else: from airflow.models import BaseOperator # type: ignore[no-redef] + from airflow.sensors.base import BaseSensorOperator # type: ignore[no-redef] -__all__ = ["AIRFLOW_V_3_0_PLUS", "BaseOperator"] +if AIRFLOW_V_3_1_PLUS: + from airflow.sdk import BaseHook +else: + from airflow.hooks.base import BaseHook # type: ignore[attr-defined,no-redef] + +__all__ = ["AIRFLOW_V_3_0_PLUS", "AIRFLOW_V_3_1_PLUS", "BaseHook", "BaseOperator", "BaseSensorOperator"]