Skip to content

Commit 1f66257

Browse files
authored
Fix MyPy errors in google.cloud.sensors (#20228)
Part of #19891
1 parent 36aa695 commit 1f66257

File tree

5 files changed

+12
-7
lines changed

5 files changed

+12
-7
lines changed

β€Žairflow/providers/google/cloud/hooks/datafusion.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -407,7 +407,7 @@ def get_pipeline_workflow(
407407
instance_url: str,
408408
pipeline_id: str,
409409
namespace: str = "default",
410-
) -> str:
410+
) -> Any:
411411
url = os.path.join(
412412
self._base_url(instance_url, namespace),
413413
quote(pipeline_name),

β€Žairflow/providers/google/cloud/sensors/bigquery_dts.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -106,7 +106,9 @@ def _normalize_state_list(self, states) -> Set[TransferState]:
106106
result = set()
107107
for state in states:
108108
if isinstance(state, str):
109-
result.add(TransferState[state.upper()])
109+
# The proto.Enum type is indexable (via MetaClass and aliased) but MyPy is not able to
110+
# infer this https://github.com/python/mypy/issues/8968
111+
result.add(TransferState[state.upper()]) # type: ignore[misc]
110112
elif isinstance(state, int):
111113
result.add(TransferState(state))
112114
elif isinstance(state, TransferState):

β€Žairflow/providers/google/cloud/sensors/datafusion.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -72,7 +72,7 @@ def __init__(
7272
expected_statuses: Set[str],
7373
instance_name: str,
7474
location: str,
75-
failure_statuses: Set[str] = None,
75+
failure_statuses: Optional[Set[str]] = None,
7676
project_id: Optional[str] = None,
7777
namespace: str = "default",
7878
gcp_conn_id: str = 'google_cloud_default',

β€Žairflow/providers/google/cloud/sensors/dataproc.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -56,7 +56,7 @@ def __init__(
5656
*,
5757
project_id: str,
5858
dataproc_job_id: str,
59-
region: str = None,
59+
region: Optional[str] = None,
6060
location: Optional[str] = None,
6161
gcp_conn_id: str = 'google_cloud_default',
6262
wait_timeout: Optional[int] = None,
@@ -79,7 +79,7 @@ def __init__(
7979
self.dataproc_job_id = dataproc_job_id
8080
self.region = region
8181
self.wait_timeout = wait_timeout
82-
self.start_sensor_time = None
82+
self.start_sensor_time: Optional[float] = None
8383

8484
def execute(self, context: Dict):
8585
self.start_sensor_time = time.monotonic()

β€Žairflow/providers/google/cloud/sensors/workflows.py

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -73,8 +73,11 @@ def __init__(
7373
):
7474
super().__init__(**kwargs)
7575

76-
self.success_states = success_states or {Execution.State.SUCCEEDED}
77-
self.failure_states = failure_states or {Execution.State.FAILED, Execution.State.CANCELLED}
76+
self.success_states = success_states or {Execution.State(Execution.State.SUCCEEDED)}
77+
self.failure_states = failure_states or {
78+
Execution.State(Execution.State.FAILED),
79+
Execution.State(Execution.State.CANCELLED),
80+
}
7881
self.workflow_id = workflow_id
7982
self.execution_id = execution_id
8083
self.location = location

0 commit comments

Comments
 (0)