yuqian90 commented on a change in pull request #10153:
URL: https://github.com/apache/airflow/pull/10153#discussion_r489949739
##########
File path: airflow/models/baseoperator.py
##########
@@ -1120,21 +1132,25 @@ def roots(self) -> List["BaseOperator"]:
"""Required by TaskMixin"""
return [self]
+ @property
+ def leaves(self) -> List["BaseOperator"]:
+ """Required by TaskMixin"""
+ return [self]
+
def _set_relatives(
self,
task_or_task_list: Union[TaskMixin, Sequence[TaskMixin]],
upstream: bool = False,
) -> None:
"""Sets relatives for the task or task list."""
-
- if isinstance(task_or_task_list, Sequence):
- task_like_object_list = task_or_task_list
- else:
- task_like_object_list = [task_or_task_list]
+ if not isinstance(task_or_task_list, Sequence):
+ task_or_task_list = [task_or_task_list]
task_list: List["BaseOperator"] = []
- for task_object in task_like_object_list:
- task_list.extend(task_object.roots)
+ for task_object in task_or_task_list:
+ task_object.update_relative(self, not upstream)
Review comment:
Hi @turbaszek I added a `update_relative` method to `TaskMixin` that
defaults to `no-op`. It's called here. `TaskGroup` overrides this method in
order to keep track of its direct upstream/downstream `TaskGroup` or
`BaseOperator` for `UI` optimization. `BaseOperator` and `XComArg` don't need
to override it.
The following is the `roots` vs `leaves` distinction that I mentioned on the
`TaskMixin` PR. I use it here so that we don't need to hardcode another `if
isinstace(task_object, TaskGroup)` here. Please see if this looks okay to you.
----------------------------------------------------------------
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.
For queries about this service, please contact Infrastructure at:
[email protected]
