From b8a6cdd2a3222d855ab91e4c8f8f0f8c98e45433 Mon Sep 17 00:00:00 2001 From: SoxMax Date: Fri, 13 Jan 2023 16:37:11 -0500 Subject: [PATCH 1/5] Update how PythonSensor returns values from python_callable --- airflow/sensors/python.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/sensors/python.py b/airflow/sensors/python.py index 615e4e20eea18..fab532509b2e6 100644 --- a/airflow/sensors/python.py +++ b/airflow/sensors/python.py @@ -71,4 +71,4 @@ def poke(self, context: Context) -> PokeReturnValue | bool: self.log.info("Poking callable: %s", str(self.python_callable)) return_value = self.python_callable(*self.op_args, **self.op_kwargs) - return PokeReturnValue(bool(return_value)) + return return_value if isinstance(return_value, PokeReturnValue) else PokeReturnValue(bool(return_value)) From 5d2b5f86b86a2f64dc85aa88fb726e01c77d8f86 Mon Sep 17 00:00:00 2001 From: SoxMax Date: Sat, 14 Jan 2023 01:24:16 +0000 Subject: [PATCH 2/5] test poke returns the xcom value --- tests/sensors/test_python.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests/sensors/test_python.py b/tests/sensors/test_python.py index 73a0ddffe45fc..e883cbb962a5c 100644 --- a/tests/sensors/test_python.py +++ b/tests/sensors/test_python.py @@ -23,6 +23,7 @@ import pytest from airflow.exceptions import AirflowSensorTimeout +from airflow.sensors.base import PokeReturnValue from airflow.sensors.python import PythonSensor from tests.operators.test_python import BasePythonTest @@ -41,6 +42,10 @@ def test_python_sensor_raise(self): with pytest.raises(ZeroDivisionError): self.run_as_task(lambda: 1 / 0) + def test_python_sensor_xcom(self): + poke_result = self.run_as_operator(fn=lambda: PokeReturnValue(True, "xcom")).poke({}) + assert poke_result.xcom_value == "xcom" + def test_python_callable_arguments_are_templatized(self): """Test PythonSensor op_args are templatized""" # Create a named tuple and ensure it is still preserved From f89a5a15aa66519d30d94f5701d3ae31a12eb88a Mon Sep 17 00:00:00 2001 From: SoxMax Date: Sun, 15 Jan 2023 18:27:23 +0000 Subject: [PATCH 3/5] update test to only run poke --- tests/sensors/test_python.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/tests/sensors/test_python.py b/tests/sensors/test_python.py index e883cbb962a5c..ec515d8dee0c7 100644 --- a/tests/sensors/test_python.py +++ b/tests/sensors/test_python.py @@ -43,7 +43,13 @@ def test_python_sensor_raise(self): self.run_as_task(lambda: 1 / 0) def test_python_sensor_xcom(self): - poke_result = self.run_as_operator(fn=lambda: PokeReturnValue(True, "xcom")).poke({}) + with self.dag: + task = self.opcls( + task_id=self.task_id, + python_callable=lambda: PokeReturnValue(True, "xcom"), + **self.default_kwargs(), + ) + poke_result = task.poke({}) assert poke_result.xcom_value == "xcom" def test_python_callable_arguments_are_templatized(self): From a2e492ac4c7b6e2ac1d4b296755cb3eeec2f07c0 Mon Sep 17 00:00:00 2001 From: SoxMax Date: Sun, 15 Jan 2023 18:28:37 +0000 Subject: [PATCH 4/5] reformat based on changes --- airflow/sensors/python.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/airflow/sensors/python.py b/airflow/sensors/python.py index fab532509b2e6..da07a2ae0c255 100644 --- a/airflow/sensors/python.py +++ b/airflow/sensors/python.py @@ -71,4 +71,6 @@ def poke(self, context: Context) -> PokeReturnValue | bool: self.log.info("Poking callable: %s", str(self.python_callable)) return_value = self.python_callable(*self.op_args, **self.op_kwargs) - return return_value if isinstance(return_value, PokeReturnValue) else PokeReturnValue(bool(return_value)) + return ( + return_value if isinstance(return_value, PokeReturnValue) else PokeReturnValue(bool(return_value)) + ) From f1160a15a769eaf4be988e691eae52b8c0883090 Mon Sep 17 00:00:00 2001 From: SoxMax Date: Mon, 16 Jan 2023 12:54:30 +0000 Subject: [PATCH 5/5] use full if rather than ternary --- airflow/sensors/python.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/airflow/sensors/python.py b/airflow/sensors/python.py index da07a2ae0c255..0a91031fd6e5f 100644 --- a/airflow/sensors/python.py +++ b/airflow/sensors/python.py @@ -71,6 +71,7 @@ def poke(self, context: Context) -> PokeReturnValue | bool: self.log.info("Poking callable: %s", str(self.python_callable)) return_value = self.python_callable(*self.op_args, **self.op_kwargs) - return ( - return_value if isinstance(return_value, PokeReturnValue) else PokeReturnValue(bool(return_value)) - ) + if isinstance(return_value, PokeReturnValue): + return return_value + else: + return PokeReturnValue(bool(return_value))