Fix MyPy errors in google.cloud.sensors (#20228)

potiuk · web-flow · commit 1f662571b213 · 2021-12-12T15:44:27.000+01:00
Part of #19891
diff --git a/airflow/providers/google/cloud/hooks/datafusion.py b/airflow/providers/google/cloud/hooks/datafusion.py
@@ -407,7 +407,7 @@ def get_pipeline_workflow(
         instance_url: str,
         pipeline_id: str,
         namespace: str = "default",
-    ) -> str:
+    ) -> Any:
         url = os.path.join(
             self._base_url(instance_url, namespace),
             quote(pipeline_name),
diff --git a/airflow/providers/google/cloud/sensors/bigquery_dts.py b/airflow/providers/google/cloud/sensors/bigquery_dts.py
@@ -106,7 +106,9 @@ def _normalize_state_list(self, states) -> Set[TransferState]:
         result = set()
         for state in states:
             if isinstance(state, str):
-                result.add(TransferState[state.upper()])
+                # The proto.Enum type is indexable (via MetaClass and aliased) but MyPy is not able to
+                # infer this https://github.com/python/mypy/issues/8968
+                result.add(TransferState[state.upper()])  # type: ignore[misc]
             elif isinstance(state, int):
                 result.add(TransferState(state))
             elif isinstance(state, TransferState):
diff --git a/airflow/providers/google/cloud/sensors/datafusion.py b/airflow/providers/google/cloud/sensors/datafusion.py
@@ -72,7 +72,7 @@ def __init__(
         expected_statuses: Set[str],
         instance_name: str,
         location: str,
-        failure_statuses: Set[str] = None,
+        failure_statuses: Optional[Set[str]] = None,
         project_id: Optional[str] = None,
         namespace: str = "default",
         gcp_conn_id: str = 'google_cloud_default',
diff --git a/airflow/providers/google/cloud/sensors/dataproc.py b/airflow/providers/google/cloud/sensors/dataproc.py
@@ -56,7 +56,7 @@ def __init__(
         *,
         project_id: str,
         dataproc_job_id: str,
-        region: str = None,
+        region: Optional[str] = None,
         location: Optional[str] = None,
         gcp_conn_id: str = 'google_cloud_default',
         wait_timeout: Optional[int] = None,
@@ -79,7 +79,7 @@ def __init__(
         self.dataproc_job_id = dataproc_job_id
         self.region = region
         self.wait_timeout = wait_timeout
-        self.start_sensor_time = None
+        self.start_sensor_time: Optional[float] = None
 
     def execute(self, context: Dict):
         self.start_sensor_time = time.monotonic()
diff --git a/airflow/providers/google/cloud/sensors/workflows.py b/airflow/providers/google/cloud/sensors/workflows.py
@@ -73,8 +73,11 @@ def __init__(
     ):
         super().__init__(**kwargs)
 
-        self.success_states = success_states or {Execution.State.SUCCEEDED}
-        self.failure_states = failure_states or {Execution.State.FAILED, Execution.State.CANCELLED}
+        self.success_states = success_states or {Execution.State(Execution.State.SUCCEEDED)}
+        self.failure_states = failure_states or {
+            Execution.State(Execution.State.FAILED),
+            Execution.State(Execution.State.CANCELLED),
+        }
         self.workflow_id = workflow_id
         self.execution_id = execution_id
         self.location = location
