ephraimbuddy commented on code in PR #32099:
URL: https://github.com/apache/airflow/pull/32099#discussion_r1242299574
##########
airflow/models/xcom_arg.py:
##########
@@ -214,6 +214,15 @@ def __enter__(self):
def __exit__(self, exc_type, exc_val, exc_tb):
SetupTeardownContext.set_work_task_roots_and_leaves()
+ @staticmethod
+ def add_task_to_context(ctx_task: Operator | PlainXComArg):
+ """Add task to context manager."""
+ if not SetupTeardownContext.active:
+ raise AirflowException("Cannot add task to context outside the
context manager.")
+ if isinstance(ctx_task, PlainXComArg):
+ ctx_task = ctx_task.operator
+ SetupTeardownContext.update_context_map(ctx_task)
Review Comment:
```suggestion
def add_task_to_context(task: Operator | PlainXComArg):
"""Add task to context manager."""
if not SetupTeardownContext.active:
raise AirflowException("Cannot add task to context outside the
context manager.")
if isinstance(task, PlainXComArg):
task = task.operator
SetupTeardownContext.update_context_map(task)
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
