fixup! AIP-72: Fix recursion bug with XComArg

Author: kaxil

airflow/sensors/base.py

@@ -51,6 +51,7 @@
 if TYPE_CHECKING:
     from sqlalchemy.orm.session import Session
 
+    from airflow.typing_compat import Self
     from airflow.utils.context import Context
 
 # As documented in https://dev.mysql.com/doc/refman/5.7/en/datetime.html.
@@ -408,7 +409,7 @@ def _get_next_poke_interval(
         self.log.info("new %s interval is %s", self.mode, new_interval)
         return new_interval
 
-    def prepare_for_execution(self) -> BaseOperator:
+    def prepare_for_execution(self) -> Self:
         task = super().prepare_for_execution()
 
         # Sensors in `poke` mode can block execution of DAGs when running

task_sdk/src/airflow/sdk/definitions/baseoperator.py

@@ -69,6 +69,7 @@
     from airflow.sdk.definitions.dag import DAG
     from airflow.sdk.definitions.taskgroup import TaskGroup
     from airflow.serialization.enums import DagAttributeTypes
+    from airflow.typing_compat import Self
     from airflow.utils.operator_resources import Resources
 
 # TODO: Task-SDK
@@ -1219,7 +1220,7 @@ def get_serialized_fields(cls):
 
         return cls.__serialized_fields
 
-    def prepare_for_execution(self) -> BaseOperator:
+    def prepare_for_execution(self) -> Self:
         """Lock task for execution to disable custom action in ``__setattr__`` and return a copy."""
         other = copy.copy(self)
         other._lock_for_execution = True

task_sdk/tests/execution_time/test_task_runner.py

@@ -147,7 +147,7 @@ def test_parse(test_dags_dir: Path, make_ti_context):
     assert isinstance(ti.task.dag, DAG)
 
 
-def test_run_basic(time_machine, mocked_parse, make_ti_context):
+def test_run_basic(time_machine, mocked_parse, make_ti_context, spy_agency):
     """Test running a basic task."""
     what = StartupDetails(
         ti=TaskInstance(id=uuid7(), task_id="hello", dag_id="super_basic_run", run_id="c", try_number=1),
@@ -163,8 +163,16 @@ def test_run_basic(time_machine, mocked_parse, make_ti_context):
         "airflow.sdk.execution_time.task_runner.SUPERVISOR_COMMS", create=True
     ) as mock_supervisor_comms:
         ti = mocked_parse(what, "super_basic_run", CustomOperator(task_id="hello"))
+
+        # Ensure that task is locked for execution
+        spy_agency.spy_on(ti.task.prepare_for_execution)
+        assert not ti.task._lock_for_execution
+
         run(ti, log=mock.MagicMock())
 
+        spy_agency.assert_spy_called(ti.task.prepare_for_execution)
+        assert ti.task._lock_for_execution
+
         mock_supervisor_comms.send_request.assert_called_once_with(
             msg=TaskState(state=TerminalTIState.SUCCESS, end_date=instant), log=mock.ANY
         )