Improve importing the module in Airflow operators package

airflow/operators/bash.py

@@ -21,14 +21,16 @@
 import shutil
 import warnings
 from functools import cached_property
-from typing import Container, Sequence
+from typing import TYPE_CHECKING, Container, Sequence
 
 from airflow.exceptions import AirflowException, AirflowSkipException
 from airflow.hooks.subprocess import SubprocessHook
 from airflow.models.baseoperator import BaseOperator
-from airflow.utils.context import Context
 from airflow.utils.operator_helpers import context_to_airflow_vars
 
+if TYPE_CHECKING:
+    from airflow.utils.context import Context
+
 
 class BashOperator(BaseOperator):
     r"""

airflow/operators/branch.py

@@ -18,11 +18,13 @@
 """Branching operators."""
 from __future__ import annotations
 
-from typing import Iterable
+from typing import TYPE_CHECKING, Iterable
 
 from airflow.models.baseoperator import BaseOperator
 from airflow.models.skipmixin import SkipMixin
-from airflow.utils.context import Context
+
+if TYPE_CHECKING:
+    from airflow.utils.context import Context
 
 
 class BaseBranchOperator(BaseOperator, SkipMixin):

airflow/operators/datetime.py

@@ -18,12 +18,14 @@
 
 import datetime
 import warnings
-from typing import Iterable
+from typing import TYPE_CHECKING, Iterable
 
 from airflow.exceptions import AirflowException, RemovedInAirflow3Warning
 from airflow.operators.branch import BaseBranchOperator
 from airflow.utils import timezone
-from airflow.utils.context import Context
+
+if TYPE_CHECKING:
+    from airflow.utils.context import Context
 
 
 class BranchDateTimeOperator(BaseBranchOperator):

airflow/operators/email.py

@@ -17,12 +17,14 @@
 # under the License.
 from __future__ import annotations
 
-from typing import Any, Sequence
+from typing import TYPE_CHECKING, Any, Sequence
 
 from airflow.models.baseoperator import BaseOperator
-from airflow.utils.context import Context
 from airflow.utils.email import send_email
 
+if TYPE_CHECKING:
+    from airflow.utils.context import Context
+
 
 class EmailOperator(BaseOperator):
     """

airflow/operators/empty.py

@@ -16,8 +16,12 @@
 # under the License.
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
+
 from airflow.models.baseoperator import BaseOperator
-from airflow.utils.context import Context
+
+if TYPE_CHECKING:
+    from airflow.utils.context import Context
 
 
 class EmptyOperator(BaseOperator):

airflow/operators/generic_transfer.py

@@ -17,11 +17,13 @@
 # under the License.
 from __future__ import annotations
 
-from typing import Sequence
+from typing import TYPE_CHECKING, Sequence
 
 from airflow.hooks.base import BaseHook
 from airflow.models import BaseOperator
-from airflow.utils.context import Context
+
+if TYPE_CHECKING:
+    from airflow.utils.context import Context
 
 
 class GenericTransfer(BaseOperator):

airflow/operators/latest_only.py

@@ -23,10 +23,10 @@
 import pendulum
 
 from airflow.operators.branch import BaseBranchOperator
-from airflow.utils.context import Context
 
 if TYPE_CHECKING:
     from airflow.models import DAG, DagRun
+    from airflow.utils.context import Context
 
 
 class LatestOnlyOperator(BaseBranchOperator):

airflow/operators/python.py

@@ -46,14 +46,16 @@
 from airflow.models.baseoperator import BaseOperator
 from airflow.models.skipmixin import SkipMixin
 from airflow.models.taskinstance import _CURRENT_CONTEXT
-from airflow.utils.context import Context, context_copy_partial, context_merge
+from airflow.utils.context import context_copy_partial, context_merge
 from airflow.utils.operator_helpers import KeywordParameters
 from airflow.utils.process_utils import execute_in_subprocess
 from airflow.utils.python_virtualenv import prepare_virtualenv, write_python_script
 
 if TYPE_CHECKING:
     from pendulum.datetime import DateTime
 
+    from airflow.utils.context import Context
+
 
 def is_venv_installed() -> bool:
     """

airflow/operators/smooth.py

@@ -17,8 +17,12 @@
 # under the License.
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
+
 from airflow.models.baseoperator import BaseOperator
-from airflow.utils.context import Context
+
+if TYPE_CHECKING:
+    from airflow.utils.context import Context
 
 
 class SmoothOperator(BaseOperator):

airflow/operators/subdag.py

@@ -23,24 +23,30 @@
 from __future__ import annotations
 
 import warnings
-from datetime import datetime
 from enum import Enum
+from typing import TYPE_CHECKING
 
 from sqlalchemy import select
-from sqlalchemy.orm.session import Session
 
 from airflow.api.common.experimental.get_task_instance import get_task_instance
 from airflow.exceptions import AirflowException, RemovedInAirflow3Warning, TaskInstanceNotFound
 from airflow.models import DagRun
-from airflow.models.dag import DAG, DagContext
+from airflow.models.dag import DagContext
 from airflow.models.pool import Pool
 from airflow.models.taskinstance import TaskInstance
 from airflow.sensors.base import BaseSensorOperator
-from airflow.utils.context import Context
 from airflow.utils.session import NEW_SESSION, create_session, provide_session
 from airflow.utils.state import DagRunState, TaskInstanceState
 from airflow.utils.types import DagRunType
 
+if TYPE_CHECKING:
+    from datetime import datetime
+
+    from sqlalchemy.orm.session import Session
+
+    from airflow.models.dag import DAG
+    from airflow.utils.context import Context
+
 
 class SkippedStatePropagationOptions(Enum):
     """Available options for skipped state propagation of subdag's tasks to parent dag tasks."""

airflow/operators/trigger_dagrun.py

@@ -34,7 +34,6 @@
 from airflow.models.xcom import XCom
 from airflow.triggers.external_task import DagStateTrigger
 from airflow.utils import timezone
-from airflow.utils.context import Context
 from airflow.utils.helpers import build_airflow_url_with_query
 from airflow.utils.session import provide_session
 from airflow.utils.state import DagRunState
@@ -48,6 +47,7 @@
     from sqlalchemy.orm.session import Session
 
     from airflow.models.taskinstancekey import TaskInstanceKey
+    from airflow.utils.context import Context
 
 
 class TriggerDagRunLink(BaseOperatorLink):

airflow/operators/weekday.py

@@ -18,14 +18,16 @@
 from __future__ import annotations
 
 import warnings
-from typing import Iterable
+from typing import TYPE_CHECKING, Iterable
 
 from airflow.exceptions import RemovedInAirflow3Warning
 from airflow.operators.branch import BaseBranchOperator
 from airflow.utils import timezone
-from airflow.utils.context import Context
 from airflow.utils.weekday import WeekDay
 
+if TYPE_CHECKING:
+    from airflow.utils.context import Context
+
 
 class BranchDayOfWeekOperator(BaseBranchOperator):
     """Branches into one of two lists of tasks depending on the current day.