yuqian90 commented on pull request #10930:
URL: https://github.com/apache/airflow/pull/10930#issuecomment-692463313
Thanks for the PR. This is a great improvement. I think there is room to
"DRY" this further. For example, we have a lot of branching logic like this in
these methods:
In `BaseOperator._set_relatives()`, we have this:
```
if isinstance(task_or_task_list, XComArg):
# otherwise we will start to iterate over xcomarg
# because of the "list" check below
# with current XComArg.__getitem__ implementation
task_list = [task_or_task_list.operator]
...
```
When we add `TaskGroup`, we'll have a few more `isinstance` checks. We might
be able to use some polymorphism to replace most of the `if isinstance` in
these methods. One idea is to introduce a `get_roots()` and `get_leaves()`
method within `TaskMixin`.
`BaseOperator` just returns itself as a list:
```
class BaseOperator(...):
def get_roots(self) -> Generator["BaseOperator", None, None]:
return [self]
def get_leaves(self) -> Generator["BaseOperator", None, None]:
return [self]
```
`XComArgs` returns the underlying operator:
```
class XComArg(...):
def get_roots(self) -> Generator["BaseOperator", None, None]:
return [self.operator]
def get_leaves(self) -> Generator["BaseOperator", None, None]:
return [self.operator]
```
`TaskGroup` has its own way of returning leaves or roots:
```
class TaskGroup(....):
def get_roots(self) -> Generator["BaseOperator", None, None]:
...
def get_leaves(self) -> Generator["BaseOperator", None, None]:
...
```
I think `get_roots()` and `get_leaves()` will be able to replace all the
`isinstance` checks in set_upstream/set_downstream of `XComArg`, `TaskGroup`,
`BaseOperator`.
What do you think?
----------------------------------------------------------------
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]
