diff --git a/airflow/sensors/python.py b/airflow/sensors/python.py index 615e4e20eea18..0a91031fd6e5f 100644 --- a/airflow/sensors/python.py +++ b/airflow/sensors/python.py @@ -71,4 +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 PokeReturnValue(bool(return_value)) + if isinstance(return_value, PokeReturnValue): + return return_value + else: + return PokeReturnValue(bool(return_value)) diff --git a/tests/sensors/test_python.py b/tests/sensors/test_python.py index 73a0ddffe45fc..ec515d8dee0c7 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,16 @@ def test_python_sensor_raise(self): with pytest.raises(ZeroDivisionError): self.run_as_task(lambda: 1 / 0) + def test_python_sensor_xcom(self): + 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): """Test PythonSensor op_args are templatized""" # Create a named tuple and ensure it is still preserved