Fix for TCH003 | stdlib imports which are only used for typechecking

airflow/api_connexion/schemas/dataset_schema.py

@@ -16,8 +16,7 @@
 # under the License.
 from __future__ import annotations
 
-from datetime import datetime
-from typing import NamedTuple
+from typing import TYPE_CHECKING, NamedTuple
 
 from marshmallow import Schema, fields
 from marshmallow_sqlalchemy import SQLAlchemySchema, auto_field
@@ -31,6 +30,8 @@
     TaskOutletDatasetReference,
 )
 
+if TYPE_CHECKING:
+    from datetime import datetime
 
 class TaskOutletDatasetReferenceSchema(SQLAlchemySchema):
     """TaskOutletDatasetReference DB schema."""

airflow/cli/commands/daemon_utils.py

@@ -17,8 +17,7 @@
 from __future__ import annotations
 
 import signal
-from argparse import Namespace
-from typing import Callable
+from typing import TYPE_CHECKING, Callable
 
 from daemon import daemon
 from daemon.pidfile import TimeoutPIDLockFile
@@ -27,6 +26,9 @@
 from airflow.utils.cli import setup_locations, setup_logging, sigint_handler, sigquit_handler
 from airflow.utils.process_utils import check_if_pidfile_process_is_running
 
+if TYPE_CHECKING:
+    from argparse import Namespace
+
 
 def run_command_with_daemon_option(
     *,

airflow/cli/commands/scheduler_command.py

@@ -19,9 +19,9 @@
 from __future__ import annotations
 
 import logging
-from argparse import Namespace
 from contextlib import contextmanager
 from multiprocessing import Process
+from typing import TYPE_CHECKING
 
 from airflow import settings
 from airflow.api_internal.internal_api_call import InternalApiConfig
@@ -36,6 +36,9 @@
 from airflow.utils.scheduler_health import serve_health_check
 from airflow.utils.usage_data_collection import usage_data_collection
 
+if TYPE_CHECKING:
+    from argparse import Namespace
+
 log = logging.getLogger(__name__)
 
 

airflow/cli/utils.py

@@ -17,11 +17,11 @@
 
 from __future__ import annotations
 
-import io
 import sys
 from typing import TYPE_CHECKING, Collection
 
 if TYPE_CHECKING:
+    import io
     from io import IOBase
 
 

airflow/datasets/manager.py

@@ -17,7 +17,6 @@
 # under the License.
 from __future__ import annotations
 
-from collections.abc import Iterable
 from typing import TYPE_CHECKING
 
 from sqlalchemy import exc, select
@@ -39,6 +38,8 @@
 from airflow.utils.session import NEW_SESSION, provide_session
 
 if TYPE_CHECKING:
+    from collections.abc import Iterable
+
     from sqlalchemy.orm.session import Session
 
     from airflow.models.dag import DagModel

airflow/io/path.py

@@ -17,7 +17,6 @@
 from __future__ import annotations
 
 import contextlib
-import os
 import shutil
 import typing
 from typing import Any, Mapping
@@ -31,6 +30,8 @@
 from airflow.io.utils.stat import stat_result
 
 if typing.TYPE_CHECKING:
+    import os
+
     from fsspec import AbstractFileSystem
 
 

airflow/models/expandinput.py

@@ -20,7 +20,6 @@
 import collections.abc
 import functools
 import operator
-from collections.abc import Sized
 from typing import TYPE_CHECKING, Any, Dict, Iterable, Mapping, NamedTuple, Sequence, Union
 
 import attr
@@ -29,6 +28,8 @@
 from airflow.utils.session import NEW_SESSION, provide_session
 
 if TYPE_CHECKING:
+    from collections.abc import Sized
+
     from sqlalchemy.orm import Session
 
     from airflow.models.operator import Operator

airflow/plugins_manager.py

@@ -27,7 +27,6 @@
 import os
 import sys
 import types
-from cgitb import Hook
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Iterable
 
@@ -42,6 +41,8 @@
 from airflow.utils.module_loading import import_string, qualname
 
 if TYPE_CHECKING:
+    from cgitb import Hook
+
     from airflow.lineage.hook import HookLineageReader
 
     try:

airflow/providers/amazon/aws/auth_manager/aws_auth_manager.py

@@ -16,7 +16,6 @@
 # under the License.
 from __future__ import annotations
 
-import argparse
 from collections import defaultdict
 from functools import cached_property
 from typing import TYPE_CHECKING, Container, Sequence, cast
@@ -53,6 +52,8 @@
     )
 
 if TYPE_CHECKING:
+    import argparse
+
     from flask_appbuilder.menu import MenuItem
 
     from airflow.auth.managers.models.base_user import BaseUser

airflow/providers/amazon/aws/executors/batch/utils.py

@@ -16,7 +16,6 @@
 # under the License.
 from __future__ import annotations
 
-import datetime
 from collections import defaultdict
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Dict, List
@@ -25,6 +24,8 @@
 from airflow.utils.state import State
 
 if TYPE_CHECKING:
+    import datetime
+
     from airflow.models.taskinstance import TaskInstanceKey
 
 CommandType = List[str]

airflow/providers/amazon/aws/executors/ecs/utils.py

@@ -23,7 +23,6 @@
 
 from __future__ import annotations
 
-import datetime
 from collections import defaultdict
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Callable, Dict, List
@@ -34,6 +33,8 @@
 from airflow.utils.state import State
 
 if TYPE_CHECKING:
+    import datetime
+
     from airflow.models.taskinstance import TaskInstanceKey
 
 CommandType = List[str]

airflow/providers/amazon/aws/log/s3_task_handler.py

@@ -17,7 +17,6 @@
 # under the License.
 from __future__ import annotations
 
-import logging
 import os
 import pathlib
 import shutil
@@ -30,6 +29,8 @@
 from airflow.utils.log.logging_mixin import LoggingMixin
 
 if TYPE_CHECKING:
+    import logging
+
     from airflow.models.taskinstance import TaskInstance
 
 

airflow/providers/databricks/operators/databricks.py

@@ -22,7 +22,6 @@
 import time
 from abc import ABC, abstractmethod
 from functools import cached_property
-from logging import Logger
 from typing import TYPE_CHECKING, Any, Sequence
 
 from deprecated import deprecated
@@ -43,6 +42,8 @@
 from airflow.providers.databricks.utils.databricks import _normalise_json_content, validate_trigger_event
 
 if TYPE_CHECKING:
+    from logging import Logger
+
     from airflow.models.taskinstancekey import TaskInstanceKey
     from airflow.utils.context import Context
     from airflow.utils.task_group import TaskGroup

airflow/providers/databricks/plugins/databricks_workflow.py

@@ -17,7 +17,6 @@
 
 from __future__ import annotations
 
-import logging
 import os
 from typing import TYPE_CHECKING, Any, cast
 from urllib.parse import unquote
@@ -46,6 +45,8 @@
 from airflow.www.views import AirflowBaseView
 
 if TYPE_CHECKING:
+    import logging
+
     from sqlalchemy.orm.session import Session
 
 

airflow/providers/fab/auth_manager/fab_auth_manager.py

@@ -17,7 +17,6 @@
 # under the License.
 from __future__ import annotations
 
-import argparse
 from functools import cached_property
 from pathlib import Path
 from typing import TYPE_CHECKING, Container
@@ -86,6 +85,8 @@
 from airflow.www.extensions.init_views import _CustomErrorRequestBodyValidator, _LazyResolver
 
 if TYPE_CHECKING:
+    import argparse
+
     from airflow.auth.managers.models.base_user import BaseUser
     from airflow.cli.cli_config import (
         CLICommand,

airflow/providers/google/cloud/hooks/dataproc.py

@@ -21,7 +21,6 @@
 
 import time
 import uuid
-from collections.abc import MutableSequence
 from typing import TYPE_CHECKING, Any, Sequence
 
 from google.api_core.client_options import ClientOptions
@@ -49,6 +48,8 @@
 from airflow.version import version as airflow_version
 
 if TYPE_CHECKING:
+    from collections.abc import MutableSequence
+
     from google.api_core.operation import Operation
     from google.api_core.operation_async import AsyncOperation
     from google.api_core.operations_v1.operations_client import OperationsClient

airflow/providers/google/cloud/operators/dataproc.py

@@ -26,7 +26,6 @@
 import time
 import uuid
 import warnings
-from collections.abc import MutableSequence
 from dataclasses import dataclass
 from datetime import datetime, timedelta
 from enum import Enum
@@ -67,6 +66,8 @@
 from airflow.utils import timezone
 
 if TYPE_CHECKING:
+    from collections.abc import MutableSequence
+
     from google.api_core import operation
     from google.api_core.retry_async import AsyncRetry
     from google.protobuf.duration_pb2 import Duration

airflow/serialization/pydantic/dag_run.py

@@ -16,7 +16,6 @@
 # under the License.
 from __future__ import annotations
 
-from datetime import datetime
 from typing import TYPE_CHECKING, Iterable, List, Optional
 
 from airflow.serialization.pydantic.dag import PydanticDag
@@ -25,6 +24,8 @@
 from airflow.utils.session import NEW_SESSION, provide_session
 
 if TYPE_CHECKING:
+    from datetime import datetime
+
     from sqlalchemy.orm import Session
 
     from airflow.jobs.scheduler_job_runner import TI

airflow/serialization/pydantic/taskinstance.py

@@ -16,7 +16,6 @@
 # under the License.
 from __future__ import annotations
 
-from datetime import datetime
 from typing import TYPE_CHECKING, Any, Iterable, Optional
 
 from typing_extensions import Annotated
@@ -46,6 +45,8 @@
 from airflow.utils.xcom import XCOM_RETURN_KEY
 
 if TYPE_CHECKING:
+    from datetime import datetime
+
     import pendulum
     from sqlalchemy.orm import Session
 

airflow/ti_deps/deps/mapped_task_upstream_dep.py

@@ -17,7 +17,6 @@
 # under the License.
 from __future__ import annotations
 
-from collections.abc import Iterator
 from typing import TYPE_CHECKING
 
 from sqlalchemy import select
@@ -27,6 +26,8 @@
 from airflow.utils.state import State, TaskInstanceState
 
 if TYPE_CHECKING:
+    from collections.abc import Iterator
+
     from sqlalchemy.orm import Session
 
     from airflow.ti_deps.dep_context import DepContext

airflow/triggers/base.py

@@ -18,11 +18,13 @@
 
 import abc
 from dataclasses import dataclass
-from datetime import timedelta
-from typing import Any, AsyncIterator
+from typing import TYPE_CHECKING, Any, AsyncIterator
 
 from airflow.utils.log.logging_mixin import LoggingMixin
 
+if TYPE_CHECKING:
+    from datetime import timedelta
+
 
 @dataclass
 class StartTriggerArgs:

airflow/utils/warnings.py

@@ -17,8 +17,11 @@
 from __future__ import annotations
 
 import warnings
-from collections.abc import Generator
 from contextlib import contextmanager
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from collections.abc import Generator
 
 
 @contextmanager