From be32277c0415f54c377564d4a7b8cc7bb92038f8 Mon Sep 17 00:00:00 2001 From: Wei Lee Date: Fri, 19 May 2023 11:18:17 +0800 Subject: [PATCH] style(mypy): fix mypy warnings --- astronomer/providers/amazon/aws/operators/sagemaker.py | 4 ++-- astronomer/providers/databricks/operators/databricks.py | 2 +- astronomer/providers/google/cloud/operators/bigquery.py | 3 ++- astronomer/providers/google/cloud/sensors/gcs.py | 2 +- 4 files changed, 6 insertions(+), 5 deletions(-) diff --git a/astronomer/providers/amazon/aws/operators/sagemaker.py b/astronomer/providers/amazon/aws/operators/sagemaker.py index 8dea544ad..038afd8fa 100644 --- a/astronomer/providers/amazon/aws/operators/sagemaker.py +++ b/astronomer/providers/amazon/aws/operators/sagemaker.py @@ -225,7 +225,7 @@ def execute(self, context: Context) -> dict[str, Any] | None: # type: ignore[ov # for bypassing mypy missing return error return None # pragma: no cover - def execute_complete(self, context: Context, event: dict[str, Any]) -> dict[str, Any]: + def execute_complete(self, context: Context, event: dict[str, Any]) -> dict[str, Any]: # type: ignore[override] """ Callback for when the trigger fires - returns immediately. Relies on trigger to throw an exception, otherwise it assumes execution was @@ -367,7 +367,7 @@ def execute(self, context: Context) -> dict[str, Any] | None: # type: ignore[ov # for bypassing mypy missing return error return None # pragma: no cover - def execute_complete(self, context: Context, event: dict[str, Any]) -> dict[str, Any]: + def execute_complete(self, context: Context, event: dict[str, Any]) -> dict[str, Any]: # type: ignore[override] """ Callback for when the trigger fires - returns immediately. Relies on trigger to throw an exception, otherwise it assumes execution was diff --git a/astronomer/providers/databricks/operators/databricks.py b/astronomer/providers/databricks/operators/databricks.py index c4af96aae..b35ffed34 100644 --- a/astronomer/providers/databricks/operators/databricks.py +++ b/astronomer/providers/databricks/operators/databricks.py @@ -222,7 +222,7 @@ def execute(self, context: Context) -> None: else: _handle_non_successful_terminal_states(run_state, run_info, hook, self.task_id) - def execute_complete(self, context: Context, event: Any = None) -> None: + def execute_complete(self, context: Context, event: Any = None) -> None: # type: ignore[override] """ Callback for when the trigger fires - returns immediately. Relies on trigger to throw an exception, otherwise it assumes execution was diff --git a/astronomer/providers/google/cloud/operators/bigquery.py b/astronomer/providers/google/cloud/operators/bigquery.py index d969bbe95..237191571 100644 --- a/astronomer/providers/google/cloud/operators/bigquery.py +++ b/astronomer/providers/google/cloud/operators/bigquery.py @@ -14,6 +14,7 @@ BigQueryValueCheckOperator, ) from google.api_core.exceptions import Conflict +from google.cloud.bigquery import CopyJob, ExtractJob, LoadJob, QueryJob from astronomer.providers.google.cloud.triggers.bigquery import ( BigQueryCheckTrigger, @@ -105,7 +106,7 @@ def execute(self, context: Context) -> Any: # noqa: D102 self._handle_job_error(job) except Conflict: # If the job already exists retrieve it - job = hook.get_job( + job: CopyJob | QueryJob | LoadJob | ExtractJob = hook.get_job( # type: ignore [no-redef] project_id=self.project_id, location=self.location, job_id=job_id, diff --git a/astronomer/providers/google/cloud/sensors/gcs.py b/astronomer/providers/google/cloud/sensors/gcs.py index 6706d5977..eeca15141 100644 --- a/astronomer/providers/google/cloud/sensors/gcs.py +++ b/astronomer/providers/google/cloud/sensors/gcs.py @@ -127,7 +127,7 @@ def __init__( ) super().__init__(**kwargs) - def execute(self, context: Context) -> None: # type: ignore[override] + def execute(self, context: Context) -> None: """Airflow runs this method on the worker and defers using the trigger.""" hook_params = {"impersonation_chain": self.impersonation_chain} if hasattr(self, "delegate_to"):