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

Part of #19891
apache · Dec 12, 2021 · 1f66257 · 1f66257
1 parent 36aa695
commit 1f66257
Show file tree

Hide file tree

Showing 5 changed files with 12 additions and 7 deletions.
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