Fix last google provider MyPy errors (#21010)

Part of #19891
apache · Jan 21, 2022 · 3728494 · 3728494
1 parent 73c0d24
commit 3728494
Show file tree

Hide file tree

Showing 5 changed files with 10 additions and 50 deletions.
diff --git a/airflow/providers/google/cloud/operators/gcs.py b/airflow/providers/google/cloud/operators/gcs.py
@@ -746,10 +746,10 @@ def execute(self, context: "Context") -> List[str]:
 
         if timespan_end is None:  # Only possible in Airflow before 2.2.
             self.log.warning("No following schedule found, setting timespan end to max %s", timespan_end)
-            timespan_end = DateTime.max
+            timespan_end = timezone.coerce_datetime(DateTime.max)
         elif timespan_start >= timespan_end:  # Airflow 2.2 sets start == end for non-perodic schedules.
             self.log.warning("DAG schedule not periodic, setting timespan end to max %s", timespan_end)
-            timespan_end = DateTime.max
+            timespan_end = timezone.coerce_datetime(DateTime.max)
 
         timespan_start = timespan_start.in_timezone(timezone.utc)
         timespan_end = timespan_end.in_timezone(timezone.utc)

diff --git a/airflow/providers/google/cloud/operators/pubsub.py b/airflow/providers/google/cloud/operators/pubsub.py
@@ -830,7 +830,7 @@ def __init__(
         subscription: str,
         max_messages: int = 5,
         ack_messages: bool = False,
-        messages_callback: Optional[Callable[[List[ReceivedMessage], Dict[str, Any]], Any]] = None,
+        messages_callback: Optional[Callable[[List[ReceivedMessage], "Context"], Any]] = None,
         gcp_conn_id: str = 'google_cloud_default',
         delegate_to: Optional[str] = None,
         impersonation_chain: Optional[Union[str, Sequence[str]]] = None,
@@ -876,7 +876,7 @@ def execute(self, context: 'Context') -> list:
     def _default_message_callback(
         self,
         pulled_messages: List[ReceivedMessage],
-        context: Dict[str, Any],
+        context: "Context",
     ) -> list:
         """
         This method can be overridden by subclasses or by `messages_callback` constructor argument.

diff --git a/airflow/providers/google/cloud/sensors/pubsub.py b/airflow/providers/google/cloud/sensors/pubsub.py
@@ -17,7 +17,7 @@
 # under the License.
 """This module contains a Google PubSub sensor."""
 import warnings
-from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Sequence, Union
+from typing import TYPE_CHECKING, Any, Callable, List, Optional, Sequence, Union
 
 from google.cloud.pubsub_v1.types import ReceivedMessage
 
@@ -104,7 +104,7 @@ def __init__(
         return_immediately: bool = True,
         ack_messages: bool = False,
         gcp_conn_id: str = 'google_cloud_default',
-        messages_callback: Optional[Callable[[List[ReceivedMessage], Dict[str, Any]], Any]] = None,
+        messages_callback: Optional[Callable[[List[ReceivedMessage], "Context"], Any]] = None,
         delegate_to: Optional[str] = None,
         project: Optional[str] = None,
         impersonation_chain: Optional[Union[str, Sequence[str]]] = None,

diff --git a/airflow/providers/google/cloud/transfers/gcs_to_local.py b/airflow/providers/google/cloud/transfers/gcs_to_local.py
@@ -89,7 +89,7 @@ def __init__(
         google_cloud_storage_conn_id: Optional[str] = None,
         delegate_to: Optional[str] = None,
         impersonation_chain: Optional[Union[str, Sequence[str]]] = None,
-        file_encoding: Optional[str] = 'utf-8',
+        file_encoding: str = 'utf-8',
         **kwargs,
     ) -> None:
         # To preserve backward compatibility

diff --git a/setup.cfg b/setup.cfg
@@ -180,49 +180,9 @@ show_error_codes = True
 [mypy-airflow.migrations.*]
 ignore_errors = True
 
-[mypy-google.cloud.tasks_v2.*]
-no_implicit_optional = False
-
-[mypy-google.cloud.automl_v1beta1.services.auto_ml.*]
-no_implicit_optional = False
-
-[mypy-google.cloud.metastore_v1.services.dataproc_metastore.*]
-no_implicit_optional = False
-
-[mypy-google.cloud.dataproc_v1.services.workflow_template_service.*]
-no_implicit_optional = False
-
-[mypy-google.cloud.dataproc_v1.services.job_controller.*]
-no_implicit_optional = False
-
-[mypy-google.cloud.dataproc_v1.services.batch_controller.*]
-no_implicit_optional = False
-
-[mypy-google.cloud.dataproc_v1.services.cluster_controller.*]
-no_implicit_optional = False
-
-[mypy-google.cloud.bigquery_datatransfer_v1.services.data_transfer_service.*]
-no_implicit_optional = False
-
-[mypy-google.cloud.devtools.cloudbuild_v1.services.cloud_build.*]
-no_implicit_optional = False
-
-[mypy-google.cloud.workflows_v1beta.services.workflows.*]
-no_implicit_optional = False
-
-[mypy-google.cloud.workflows.executions_v1beta.services.executions.*]
-no_implicit_optional = False
-
-[mypy-google.cloud.datacatalog_v1.services.data_catalog.*]
-no_implicit_optional = False
-
-[mypy-google.cloud.redis_v1.services.cloud_redis.*]
-no_implicit_optional = False
-
-[mypy-google.cloud.kms_v1.services.key_management_service.*]
-no_implicit_optional = False
-
-[mypy-google.cloud.oslogin_v1.services.os_login_service.*]
+# Let's assume all google.cloud packages have no implicit optional
+# Most of them don't but even if they do, it does not matter
+[mypy-google.cloud.*]
 no_implicit_optional = False
 
 [isort]