From 45983b600b5fd5c27650f2ac9b7bf30cdb8de593 Mon Sep 17 00:00:00 2001 From: Wei Lee Date: Tue, 4 Jun 2024 05:07:20 +0800 Subject: [PATCH] style: fix mypy warning (#1541) --- .../snowflake/example_dags/example_snowflake_sensor.py | 2 +- astronomer/providers/snowflake/sensors/snowflake.py | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/astronomer/providers/snowflake/example_dags/example_snowflake_sensor.py b/astronomer/providers/snowflake/example_dags/example_snowflake_sensor.py index 6a9717a52..53ba7fc2c 100644 --- a/astronomer/providers/snowflake/example_dags/example_snowflake_sensor.py +++ b/astronomer/providers/snowflake/example_dags/example_snowflake_sensor.py @@ -67,7 +67,7 @@ sql=SNOWFLAKE_SLACK_SQL, poke_interval=POKE_INTERVAL, timeout=TASK_TIMEOUT * 60, - success=lambda result: True, # type: ignore[arg-type] + success=lambda result: True, ) snowflake_op_sql_str >> snowflake_op_with_params >> snowflake_op_sql_sensor >> snowflake_with_callable diff --git a/astronomer/providers/snowflake/sensors/snowflake.py b/astronomer/providers/snowflake/sensors/snowflake.py index b40d0a743..f5f5106b9 100644 --- a/astronomer/providers/snowflake/sensors/snowflake.py +++ b/astronomer/providers/snowflake/sensors/snowflake.py @@ -1,7 +1,7 @@ from __future__ import annotations from datetime import timedelta -from typing import Any, Sequence +from typing import Any, Callable, Sequence from airflow.exceptions import AirflowException from airflow.providers.common.sql.sensors.sql import SqlSensor @@ -48,8 +48,8 @@ def __init__( snowflake_conn_id: str, sql: str, parameters: dict[str, Any] | None = None, - success: str | None = None, - failure: str | None = None, + success: Callable[[Any], bool] | None = None, + failure: Callable[[Any], bool] | None = None, fail_on_empty: bool = False, hook_params: dict[str, Any] | None = None, **kwargs: Any,