From 49a0c0bbaf85d4c93205ea9f297f4d43c1e832eb Mon Sep 17 00:00:00 2001 From: morrme Date: Sat, 11 Jul 2020 00:22:21 -0500 Subject: [PATCH 1/4] improve type hinting for celery provider Per issue #9708 --- airflow/providers/celery/sensors/celery_queue.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/airflow/providers/celery/sensors/celery_queue.py b/airflow/providers/celery/sensors/celery_queue.py index d3ca42794c300..67ab3f60d7b85 100644 --- a/airflow/providers/celery/sensors/celery_queue.py +++ b/airflow/providers/celery/sensors/celery_queue.py @@ -17,7 +17,7 @@ # under the License. from celery.app import control - +from typing import Optional, Dict, Any from airflow.sensors.base_sensor_operator import BaseSensorOperator from airflow.utils.decorators import apply_defaults @@ -36,16 +36,16 @@ class CeleryQueueSensor(BaseSensorOperator): @apply_defaults def __init__( self, - celery_queue, - target_task_id=None, + celery_queue: str, + target_task_id: Optional[str]=None, *args, - **kwargs): + **kwargs) -> None: super().__init__(*args, **kwargs) self.celery_queue = celery_queue self.target_task_id = target_task_id - def _check_task_id(self, context): + def _check_task_id(self, context: Dict[str, Any]) -> bool: """ Gets the returned Celery result from the Airflow task ID provided to the sensor, and returns True if the @@ -60,7 +60,7 @@ def _check_task_id(self, context): celery_result = ti.xcom_pull(task_ids=self.target_task_id) return celery_result.ready() - def poke(self, context): + def poke(self, context: Dict[str,Any]) -> bool: if self.target_task_id: return self._check_task_id(context) From 4b2ec4722f1748a07e8deaf4d81e65e9eb829169 Mon Sep 17 00:00:00 2001 From: morrme Date: Sat, 11 Jul 2020 00:41:00 -0500 Subject: [PATCH 2/4] fix linter errors --- airflow/providers/celery/sensors/celery_queue.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/airflow/providers/celery/sensors/celery_queue.py b/airflow/providers/celery/sensors/celery_queue.py index 67ab3f60d7b85..02abf2c4704fb 100644 --- a/airflow/providers/celery/sensors/celery_queue.py +++ b/airflow/providers/celery/sensors/celery_queue.py @@ -37,7 +37,7 @@ class CeleryQueueSensor(BaseSensorOperator): def __init__( self, celery_queue: str, - target_task_id: Optional[str]=None, + target_task_id: Optional[str] = None, *args, **kwargs) -> None: @@ -60,7 +60,7 @@ def _check_task_id(self, context: Dict[str, Any]) -> bool: celery_result = ti.xcom_pull(task_ids=self.target_task_id) return celery_result.ready() - def poke(self, context: Dict[str,Any]) -> bool: + def poke(self, context: Dict[str, Any]) -> bool: if self.target_task_id: return self._check_task_id(context) From 01bcf33b217ca3c8a035578213905ebcb1861678 Mon Sep 17 00:00:00 2001 From: morrme Date: Sat, 11 Jul 2020 00:57:15 -0500 Subject: [PATCH 3/4] correct import sort --- airflow/providers/celery/sensors/celery_queue.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/providers/celery/sensors/celery_queue.py b/airflow/providers/celery/sensors/celery_queue.py index 02abf2c4704fb..0d089001ff722 100644 --- a/airflow/providers/celery/sensors/celery_queue.py +++ b/airflow/providers/celery/sensors/celery_queue.py @@ -17,7 +17,7 @@ # under the License. from celery.app import control -from typing import Optional, Dict, Any +from typing import Any, Dict, Optional from airflow.sensors.base_sensor_operator import BaseSensorOperator from airflow.utils.decorators import apply_defaults From a77ef6c167eaf52a8f99903eeec1df594dc156fd Mon Sep 17 00:00:00 2001 From: morrme Date: Sat, 11 Jul 2020 01:14:55 -0500 Subject: [PATCH 4/4] Update celery_queue.py --- airflow/providers/celery/sensors/celery_queue.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/airflow/providers/celery/sensors/celery_queue.py b/airflow/providers/celery/sensors/celery_queue.py index 0d089001ff722..f9b09e57d33b6 100644 --- a/airflow/providers/celery/sensors/celery_queue.py +++ b/airflow/providers/celery/sensors/celery_queue.py @@ -16,8 +16,10 @@ # specific language governing permissions and limitations # under the License. -from celery.app import control from typing import Any, Dict, Optional + +from celery.app import control + from airflow.sensors.base_sensor_operator import BaseSensorOperator from airflow.utils.decorators import apply_defaults