Convert apache providers and its tests to __future__.annotations

airflow/providers/apache/beam/hooks/beam.py

@@ -16,6 +16,8 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains a Apache Beam Hook."""
+from __future__ import annotations
+
 import json
 import os
 import select
@@ -24,7 +26,7 @@
 import subprocess
 import textwrap
 from tempfile import TemporaryDirectory
-from typing import Callable, List, Optional
+from typing import Callable
 
 from airflow.exceptions import AirflowConfigException, AirflowException
 from airflow.hooks.base import BaseHook
@@ -50,7 +52,7 @@ class BeamRunnerType:
     Twister2Runner = "Twister2Runner"
 
 
-def beam_options_to_args(options: dict) -> List[str]:
+def beam_options_to_args(options: dict) -> list[str]:
     """
     Returns a formatted pipeline options from a dictionary of arguments
 
@@ -65,7 +67,7 @@ def beam_options_to_args(options: dict) -> List[str]:
     if not options:
         return []
 
-    args: List[str] = []
+    args: list[str] = []
     for attr, value in options.items():
         if value is None or (isinstance(value, bool) and value):
             args.append(f"--{attr}")
@@ -88,14 +90,14 @@ class BeamCommandRunner(LoggingMixin):
 
     def __init__(
         self,
-        cmd: List[str],
-        process_line_callback: Optional[Callable[[str], None]] = None,
-        working_directory: Optional[str] = None,
+        cmd: list[str],
+        process_line_callback: Callable[[str], None] | None = None,
+        working_directory: str | None = None,
     ) -> None:
         super().__init__()
         self.log.info("Running command: %s", " ".join(shlex.quote(c) for c in cmd))
         self.process_line_callback = process_line_callback
-        self.job_id: Optional[str] = None
+        self.job_id: str | None = None
 
         self._proc = subprocess.Popen(
             cmd,
@@ -173,9 +175,9 @@ def __init__(
     def _start_pipeline(
         self,
         variables: dict,
-        command_prefix: List[str],
-        process_line_callback: Optional[Callable[[str], None]] = None,
-        working_directory: Optional[str] = None,
+        command_prefix: list[str],
+        process_line_callback: Callable[[str], None] | None = None,
+        working_directory: str | None = None,
     ) -> None:
         cmd = command_prefix + [
             f"--runner={self.runner}",
@@ -193,11 +195,11 @@ def start_python_pipeline(
         self,
         variables: dict,
         py_file: str,
-        py_options: List[str],
+        py_options: list[str],
         py_interpreter: str = "python3",
-        py_requirements: Optional[List[str]] = None,
+        py_requirements: list[str] | None = None,
         py_system_site_packages: bool = False,
-        process_line_callback: Optional[Callable[[str], None]] = None,
+        process_line_callback: Callable[[str], None] | None = None,
     ):
         """
         Starts Apache Beam python pipeline.
@@ -266,8 +268,8 @@ def start_java_pipeline(
         self,
         variables: dict,
         jar: str,
-        job_class: Optional[str] = None,
-        process_line_callback: Optional[Callable[[str], None]] = None,
+        job_class: str | None = None,
+        process_line_callback: Callable[[str], None] | None = None,
     ) -> None:
         """
         Starts Apache Beam Java pipeline.
@@ -292,7 +294,7 @@ def start_go_pipeline(
         self,
         variables: dict,
         go_file: str,
-        process_line_callback: Optional[Callable[[str], None]] = None,
+        process_line_callback: Callable[[str], None] | None = None,
         should_init_module: bool = False,
     ) -> None:
         """

airflow/providers/apache/beam/operators/beam.py

@@ -16,11 +16,13 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module contains Apache Beam operators."""
+from __future__ import annotations
+
 import copy
 import tempfile
 from abc import ABC, ABCMeta
 from contextlib import ExitStack
-from typing import TYPE_CHECKING, Callable, List, Optional, Sequence, Tuple, Union
+from typing import TYPE_CHECKING, Callable, Sequence
 
 from airflow import AirflowException
 from airflow.models import BaseOperator
@@ -47,17 +49,17 @@ class BeamDataflowMixin(metaclass=ABCMeta):
     :class:`~airflow.providers.apache.beam.operators.beam.BeamRunGoPipelineOperator`.
     """
 
-    dataflow_hook: Optional[DataflowHook]
+    dataflow_hook: DataflowHook | None
     dataflow_config: DataflowConfiguration
     gcp_conn_id: str
-    delegate_to: Optional[str]
+    delegate_to: str | None
     dataflow_support_impersonation: bool = True
 
     def _set_dataflow(
         self,
         pipeline_options: dict,
-        job_name_variable_key: Optional[str] = None,
-    ) -> Tuple[str, dict, Callable[[str], None]]:
+        job_name_variable_key: str | None = None,
+    ) -> tuple[str, dict, Callable[[str], None]]:
         self.dataflow_hook = self.__set_dataflow_hook()
         self.dataflow_config.project_id = self.dataflow_config.project_id or self.dataflow_hook.project_id
         dataflow_job_name = self.__get_dataflow_job_name()
@@ -85,7 +87,7 @@ def __get_dataflow_job_name(self) -> str:
         )
 
     def __get_dataflow_pipeline_options(
-        self, pipeline_options: dict, job_name: str, job_name_key: Optional[str] = None
+        self, pipeline_options: dict, job_name: str, job_name_key: str | None = None
     ) -> dict:
         pipeline_options = copy.deepcopy(pipeline_options)
         if job_name_key is not None:
@@ -151,11 +153,11 @@ def __init__(
         self,
         *,
         runner: str = "DirectRunner",
-        default_pipeline_options: Optional[dict] = None,
-        pipeline_options: Optional[dict] = None,
+        default_pipeline_options: dict | None = None,
+        pipeline_options: dict | None = None,
         gcp_conn_id: str = "google_cloud_default",
-        delegate_to: Optional[str] = None,
-        dataflow_config: Optional[Union[DataflowConfiguration, dict]] = None,
+        delegate_to: str | None = None,
+        dataflow_config: DataflowConfiguration | dict | None = None,
         **kwargs,
     ) -> None:
         super().__init__(**kwargs)
@@ -168,9 +170,9 @@ def __init__(
             self.dataflow_config = DataflowConfiguration(**dataflow_config)
         else:
             self.dataflow_config = dataflow_config or DataflowConfiguration()
-        self.beam_hook: Optional[BeamHook] = None
-        self.dataflow_hook: Optional[DataflowHook] = None
-        self.dataflow_job_id: Optional[str] = None
+        self.beam_hook: BeamHook | None = None
+        self.dataflow_hook: DataflowHook | None = None
+        self.dataflow_job_id: str | None = None
 
         if self.dataflow_config and self.runner.lower() != BeamRunnerType.DataflowRunner.lower():
             self.log.warning(
@@ -180,13 +182,13 @@ def __init__(
     def _init_pipeline_options(
         self,
         format_pipeline_options: bool = False,
-        job_name_variable_key: Optional[str] = None,
-    ) -> Tuple[bool, Optional[str], dict, Optional[Callable[[str], None]]]:
+        job_name_variable_key: str | None = None,
+    ) -> tuple[bool, str | None, dict, Callable[[str], None] | None]:
         self.beam_hook = BeamHook(runner=self.runner)
         pipeline_options = self.default_pipeline_options.copy()
-        process_line_callback: Optional[Callable[[str], None]] = None
+        process_line_callback: Callable[[str], None] | None = None
         is_dataflow = self.runner.lower() == BeamRunnerType.DataflowRunner.lower()
-        dataflow_job_name: Optional[str] = None
+        dataflow_job_name: str | None = None
         if is_dataflow:
             dataflow_job_name, pipeline_options, process_line_callback = self._set_dataflow(
                 pipeline_options=pipeline_options,
@@ -255,15 +257,15 @@ def __init__(
         *,
         py_file: str,
         runner: str = "DirectRunner",
-        default_pipeline_options: Optional[dict] = None,
-        pipeline_options: Optional[dict] = None,
+        default_pipeline_options: dict | None = None,
+        pipeline_options: dict | None = None,
         py_interpreter: str = "python3",
-        py_options: Optional[List[str]] = None,
-        py_requirements: Optional[List[str]] = None,
+        py_options: list[str] | None = None,
+        py_requirements: list[str] | None = None,
         py_system_site_packages: bool = False,
         gcp_conn_id: str = "google_cloud_default",
-        delegate_to: Optional[str] = None,
-        dataflow_config: Optional[Union[DataflowConfiguration, dict]] = None,
+        delegate_to: str | None = None,
+        dataflow_config: DataflowConfiguration | dict | None = None,
         **kwargs,
     ) -> None:
         super().__init__(
@@ -285,7 +287,7 @@ def __init__(
             {"airflow-version": "v" + version.replace(".", "-").replace("+", "-")}
         )
 
-    def execute(self, context: 'Context'):
+    def execute(self, context: Context):
         """Execute the Apache Beam Pipeline."""
         (
             is_dataflow,
@@ -396,12 +398,12 @@ def __init__(
         *,
         jar: str,
         runner: str = "DirectRunner",
-        job_class: Optional[str] = None,
-        default_pipeline_options: Optional[dict] = None,
-        pipeline_options: Optional[dict] = None,
+        job_class: str | None = None,
+        default_pipeline_options: dict | None = None,
+        pipeline_options: dict | None = None,
         gcp_conn_id: str = "google_cloud_default",
-        delegate_to: Optional[str] = None,
-        dataflow_config: Optional[Union[DataflowConfiguration, dict]] = None,
+        delegate_to: str | None = None,
+        dataflow_config: DataflowConfiguration | dict | None = None,
         **kwargs,
     ) -> None:
         super().__init__(
@@ -416,7 +418,7 @@ def __init__(
         self.jar = jar
         self.job_class = job_class
 
-    def execute(self, context: 'Context'):
+    def execute(self, context: Context):
         """Execute the Apache Beam Pipeline."""
         (
             is_dataflow,
@@ -537,11 +539,11 @@ def __init__(
         *,
         go_file: str,
         runner: str = "DirectRunner",
-        default_pipeline_options: Optional[dict] = None,
-        pipeline_options: Optional[dict] = None,
+        default_pipeline_options: dict | None = None,
+        pipeline_options: dict | None = None,
         gcp_conn_id: str = "google_cloud_default",
-        delegate_to: Optional[str] = None,
-        dataflow_config: Optional[Union[DataflowConfiguration, dict]] = None,
+        delegate_to: str | None = None,
+        dataflow_config: DataflowConfiguration | dict | None = None,
         **kwargs,
     ) -> None:
         super().__init__(
@@ -567,7 +569,7 @@ def __init__(
             {"airflow-version": "v" + version.replace(".", "-").replace("+", "-")}
         )
 
-    def execute(self, context: 'Context'):
+    def execute(self, context: Context):
         """Execute the Apache Beam Pipeline."""
         (
             is_dataflow,

airflow/providers/apache/cassandra/hooks/cassandra.py

@@ -15,10 +15,10 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-
 """This module contains hook to integrate with Apache Cassandra."""
+from __future__ import annotations
 
-from typing import Any, Dict, Union
+from typing import Any, Union
 
 from cassandra.auth import PlainTextAuthProvider
 from cassandra.cluster import Cluster, Session
@@ -141,7 +141,7 @@ def shutdown_cluster(self) -> None:
             self.cluster.shutdown()
 
     @staticmethod
-    def get_lb_policy(policy_name: str, policy_args: Dict[str, Any]) -> Policy:
+    def get_lb_policy(policy_name: str, policy_args: dict[str, Any]) -> Policy:
         """
         Creates load balancing policy.
 
@@ -188,7 +188,7 @@ def table_exists(self, table: str) -> bool:
         cluster_metadata = self.get_conn().cluster.metadata
         return keyspace in cluster_metadata.keyspaces and table in cluster_metadata.keyspaces[keyspace].tables
 
-    def record_exists(self, table: str, keys: Dict[str, str]) -> bool:
+    def record_exists(self, table: str, keys: dict[str, str]) -> bool:
         """
         Checks if a record exists in Cassandra
 

airflow/providers/apache/cassandra/sensors/record.py

@@ -19,8 +19,9 @@
 This module contains sensor that check the existence
 of a record in a Cassandra cluster.
 """
+from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Dict, Sequence
+from typing import TYPE_CHECKING, Any, Sequence
 
 from airflow.providers.apache.cassandra.hooks.cassandra import CassandraHook
 from airflow.sensors.base import BaseSensorOperator
@@ -58,7 +59,7 @@ class CassandraRecordSensor(BaseSensorOperator):
     def __init__(
         self,
         *,
-        keys: Dict[str, str],
+        keys: dict[str, str],
         table: str,
         cassandra_conn_id: str = CassandraHook.default_conn_name,
         **kwargs: Any,
@@ -68,7 +69,7 @@ def __init__(
         self.table = table
         self.keys = keys
 
-    def poke(self, context: "Context") -> bool:
+    def poke(self, context: Context) -> bool:
         self.log.info('Sensor check existence of record: %s', self.keys)
         hook = CassandraHook(self.cassandra_conn_id)
         return hook.record_exists(self.table, self.keys)

airflow/providers/apache/cassandra/sensors/table.py

@@ -15,11 +15,11 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-
 """
 This module contains sensor that check the existence
 of a table in a Cassandra cluster.
 """
+from __future__ import annotations
 
 from typing import TYPE_CHECKING, Any, Sequence
 
@@ -65,7 +65,7 @@ def __init__(
         self.cassandra_conn_id = cassandra_conn_id
         self.table = table
 
-    def poke(self, context: "Context") -> bool:
+    def poke(self, context: Context) -> bool:
         self.log.info('Sensor check existence of table: %s', self.table)
         hook = CassandraHook(self.cassandra_conn_id)
         return hook.table_exists(self.table)

airflow/providers/apache/drill/hooks/drill.py

@@ -15,8 +15,9 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+from __future__ import annotations
 
-from typing import Any, Iterable, Optional, Tuple
+from typing import Any, Iterable
 
 from sqlalchemy import create_engine
 from sqlalchemy.engine import Connection
@@ -80,8 +81,8 @@ def set_autocommit(self, conn: Connection, autocommit: bool) -> NotImplementedEr
     def insert_rows(
         self,
         table: str,
-        rows: Iterable[Tuple[str]],
-        target_fields: Optional[Iterable[str]] = None,
+        rows: Iterable[tuple[str]],
+        target_fields: Iterable[str] | None = None,
         commit_every: int = 1000,
         replace: bool = False,
         **kwargs: Any,