Improve modules import in cncf.kubernetes probvider by move some of t…

…hem into a type-checking block (#33781)

airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py

@@ -23,7 +23,6 @@
 """
 from __future__ import annotations
 
-import argparse
 import json
 import logging
 import multiprocessing
@@ -35,7 +34,6 @@
 from typing import TYPE_CHECKING, Any, Sequence
 
 from sqlalchemy import select, update
-from sqlalchemy.orm import Session
 
 from airflow.exceptions import AirflowException
 
@@ -83,8 +81,11 @@
 from airflow.utils.state import TaskInstanceState
 
 if TYPE_CHECKING:
+    import argparse
+
     from kubernetes import client
     from kubernetes.client import models as k8s
+    from sqlalchemy.orm import Session
 
     from airflow.executors.base_executor import CommandType
     from airflow.models.taskinstance import TaskInstance

airflow/providers/cncf/kubernetes/executors/kubernetes_executor_utils.py

@@ -23,7 +23,6 @@
 from typing import TYPE_CHECKING, Any, Generic, TypeVar
 
 from kubernetes import client, watch
-from kubernetes.client import Configuration, models as k8s
 from kubernetes.client.rest import ApiException
 from urllib3.exceptions import ReadTimeoutError
 
@@ -51,6 +50,8 @@
     )
 
 if TYPE_CHECKING:
+    from kubernetes.client import Configuration, models as k8s
+
     from airflow.providers.cncf.kubernetes.executors.kubernetes_executor_types import (
         KubernetesJobType,
         KubernetesResultsType,

airflow/providers/cncf/kubernetes/executors/local_kubernetes_executor.py

@@ -19,15 +19,15 @@
 
 from typing import TYPE_CHECKING, Sequence
 
-from airflow.callbacks.base_callback_sink import BaseCallbackSink
-from airflow.callbacks.callback_requests import CallbackRequest
 from airflow.configuration import conf
-from airflow.executors.local_executor import LocalExecutor
 from airflow.providers.cncf.kubernetes.executors.kubernetes_executor import KubernetesExecutor
 from airflow.utils.log.logging_mixin import LoggingMixin
 
 if TYPE_CHECKING:
+    from airflow.callbacks.base_callback_sink import BaseCallbackSink
+    from airflow.callbacks.callback_requests import CallbackRequest
     from airflow.executors.base_executor import CommandType, EventBufferValueType, QueuedTaskInstanceType
+    from airflow.executors.local_executor import LocalExecutor
     from airflow.models.taskinstance import SimpleTaskInstance, TaskInstance, TaskInstanceKey
 
 

airflow/providers/cncf/kubernetes/hooks/kubernetes.py

@@ -24,7 +24,6 @@
 
 from asgiref.sync import sync_to_async
 from kubernetes import client, config, watch
-from kubernetes.client.models import V1Pod
 from kubernetes.config import ConfigException
 from kubernetes_asyncio import client as async_client, config as async_config
 from urllib3.exceptions import HTTPError
@@ -36,6 +35,9 @@
 from airflow.providers.cncf.kubernetes.utils.pod_manager import PodOperatorHookProtocol
 from airflow.utils import yaml
 
+if TYPE_CHECKING:
+    from kubernetes.client.models import V1Pod
+
 LOADING_KUBE_CONFIG_FILE_RESOURCE = "Loading Kubernetes configuration file kube_config from {}..."
 
 

airflow/providers/cncf/kubernetes/k8s_model.py

@@ -19,8 +19,10 @@
 
 from abc import ABC, abstractmethod
 from functools import reduce
+from typing import TYPE_CHECKING
 
-from kubernetes.client import models as k8s
+if TYPE_CHECKING:
+    from kubernetes.client import models as k8s
 
 
 class K8SModel(ABC):

airflow/providers/cncf/kubernetes/operators/pod.py

@@ -52,7 +52,6 @@
 )
 from airflow.providers.cncf.kubernetes.hooks.kubernetes import KubernetesHook
 from airflow.providers.cncf.kubernetes.pod_generator import PodGenerator
-from airflow.providers.cncf.kubernetes.secret import Secret
 from airflow.providers.cncf.kubernetes.triggers.pod import KubernetesPodTrigger
 from airflow.providers.cncf.kubernetes.utils import xcom_sidecar  # type: ignore[attr-defined]
 from airflow.providers.cncf.kubernetes.utils.pod_manager import (
@@ -65,7 +64,6 @@
     get_container_termination_message,
 )
 from airflow.settings import pod_mutation_hook
-from airflow.typing_compat import Literal
 from airflow.utils import yaml
 from airflow.utils.helpers import prune_dict, validate_key
 from airflow.utils.timezone import utcnow
@@ -74,6 +72,8 @@
 if TYPE_CHECKING:
     import jinja2
 
+    from airflow.providers.cncf.kubernetes.secret import Secret
+    from airflow.typing_compat import Literal
     from airflow.utils.context import Context
 
 alphanum_lower = string.ascii_lowercase + string.digits

airflow/providers/cncf/kubernetes/operators/resource.py

@@ -19,15 +19,18 @@
 from __future__ import annotations
 
 from functools import cached_property
+from typing import TYPE_CHECKING
 
 import yaml
-from kubernetes.client import ApiClient
 from kubernetes.utils import create_from_yaml
 
 from airflow.models import BaseOperator
 from airflow.providers.cncf.kubernetes.hooks.kubernetes import KubernetesHook
 from airflow.providers.cncf.kubernetes.utils.delete_from import delete_from_yaml
 
+if TYPE_CHECKING:
+    from kubernetes.client import ApiClient
+
 __all__ = ["KubernetesCreateResourceOperator", "KubernetesDeleteResourceOperator"]
 
 

airflow/providers/cncf/kubernetes/pod_generator.py

@@ -25,11 +25,11 @@
 from __future__ import annotations
 
 import copy
-import datetime
 import logging
 import os
 import warnings
 from functools import reduce
+from typing import TYPE_CHECKING
 
 import re2
 from dateutil import parser
@@ -55,6 +55,9 @@
 from airflow.utils import yaml
 from airflow.version import version as airflow_version
 
+if TYPE_CHECKING:
+    import datetime
+
 log = logging.getLogger(__name__)
 
 MAX_LABEL_LEN = 63

airflow/providers/cncf/kubernetes/pod_launcher_deprecated.py

@@ -22,11 +22,11 @@
 import time
 import warnings
 from datetime import datetime as dt
+from typing import TYPE_CHECKING
 
 import pendulum
 import tenacity
 from kubernetes import client, watch
-from kubernetes.client.models.v1_pod import V1Pod
 from kubernetes.client.rest import ApiException
 from kubernetes.stream import stream as kubernetes_stream
 from requests.exceptions import HTTPError
@@ -38,6 +38,9 @@
 from airflow.utils.log.logging_mixin import LoggingMixin
 from airflow.utils.state import State
 
+if TYPE_CHECKING:
+    from kubernetes.client.models.v1_pod import V1Pod
+
 warnings.warn(
     """
     Please use :mod: Please use `airflow.providers.cncf.kubernetes.utils.pod_manager`

airflow/providers/cncf/kubernetes/template_rendering.py

@@ -17,18 +17,22 @@
 
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
+
 from jinja2 import TemplateAssertionError, UndefinedError
 from kubernetes.client.api_client import ApiClient
 
 from airflow import AirflowException
-from airflow.models.taskinstance import TaskInstance
 from airflow.providers.cncf.kubernetes.kube_config import KubeConfig
 from airflow.providers.cncf.kubernetes.kubernetes_helper_functions import (
     create_pod_id,
 )
 from airflow.providers.cncf.kubernetes.pod_generator import PodGenerator
 from airflow.utils.session import NEW_SESSION, provide_session
 
+if TYPE_CHECKING:
+    from airflow.models.taskinstance import TaskInstance
+
 
 def render_k8s_pod_yaml(task_instance: TaskInstance) -> dict | None:
     """Render k8s pod yaml."""

airflow/providers/cncf/kubernetes/triggers/pod.py

@@ -21,15 +21,16 @@
 import warnings
 from asyncio import CancelledError
 from enum import Enum
-from typing import Any, AsyncIterator
-
-from kubernetes_asyncio.client.models import V1Pod
+from typing import TYPE_CHECKING, Any, AsyncIterator
 
 from airflow.exceptions import AirflowProviderDeprecationWarning
 from airflow.providers.cncf.kubernetes.hooks.kubernetes import AsyncKubernetesHook
 from airflow.providers.cncf.kubernetes.utils.pod_manager import OnFinishAction, PodPhase
 from airflow.triggers.base import BaseTrigger, TriggerEvent
 
+if TYPE_CHECKING:
+    from kubernetes_asyncio.client.models import V1Pod
+
 
 class ContainerState(str, Enum):
     """

airflow/providers/cncf/kubernetes/utils/delete_from.py

@@ -20,9 +20,12 @@
 from __future__ import annotations
 
 import re
+from typing import TYPE_CHECKING
 
 from kubernetes import client
-from kubernetes.client import ApiClient
+
+if TYPE_CHECKING:
+    from kubernetes.client import ApiClient
 
 DEFAULT_DELETION_BODY = client.V1DeleteOptions(
     propagation_policy="Background",

airflow/providers/cncf/kubernetes/utils/pod_manager.py

@@ -33,15 +33,12 @@
 import pendulum
 import tenacity
 from kubernetes import client, watch
-from kubernetes.client.models.v1_container_status import V1ContainerStatus
-from kubernetes.client.models.v1_pod import V1Pod
 from kubernetes.client.rest import ApiException
 from kubernetes.stream import stream as kubernetes_stream
 from pendulum import DateTime
 from pendulum.parsing.exceptions import ParserError
 from tenacity import before_log
 from urllib3.exceptions import HTTPError as BaseHTTPError
-from urllib3.response import HTTPResponse
 
 from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
 from airflow.providers.cncf.kubernetes.pod_generator import PodDefaults
@@ -51,6 +48,9 @@
 
 if TYPE_CHECKING:
     from kubernetes.client.models.core_v1_event_list import CoreV1EventList
+    from kubernetes.client.models.v1_container_status import V1ContainerStatus
+    from kubernetes.client.models.v1_pod import V1Pod
+    from urllib3.response import HTTPResponse
 
 
 class PodLaunchFailedException(AirflowException):