ashb commented on a change in pull request #19965:
URL: https://github.com/apache/airflow/pull/19965#discussion_r771510404



##########
File path: airflow/models/taskmixin.py
##########
@@ -88,3 +99,167 @@ def __init_subclass__(cls) -> None:
             stacklevel=2,
         )
         return super().__init_subclass__()
+
+
+class DAGNode(DependencyMixin):
+    """
+    A base class for a node in the graph of a workflow -- an Operator or a 
Task Group, either mapped or
+    unmapped.
+    """
+
+    dag: Optional["DAG"] = None
+
+    @property
+    @abstractmethod
+    def node_id(self) -> str:
+        raise NotImplementedError()
+
+    task_group: Optional["TaskGroup"]
+    """The task_group that contains this node"""
+
+    start_date: Optional[pendulum.DateTime]
+    end_date: Optional[pendulum.DateTime]
+
+    def has_dag(self) -> bool:
+        return self.dag is not None
+
+    @property
+    @abstractmethod
+    def upstream_task_ids(self) -> Set[str]:
+        raise NotImplementedError()
+
+    @property
+    @abstractmethod
+    def downstream_task_ids(self) -> Set[str]:
+        raise NotImplementedError()
+
+    @property
+    def log(self) -> "Logger":
+        raise NotImplementedError()
+
+    @property
+    @abstractmethod
+    def roots(self) -> Sequence["DAGNode"]:
+        raise NotImplementedError()
+
+    @property
+    @abstractmethod
+    def leaves(self) -> Sequence["DAGNode"]:
+        raise NotImplementedError()
+
+    def _set_relatives(
+        self,
+        task_or_task_list: Union[DependencyMixin, Sequence[DependencyMixin]],
+        upstream: bool = False,
+        edge_modifier: Optional["EdgeModifier"] = None,
+    ) -> None:
+        """Sets relatives for the task or task list."""
+        from airflow.models.baseoperator import BaseOperator, MappedOperator
+
+        if not isinstance(task_or_task_list, Sequence):
+            task_or_task_list = [task_or_task_list]
+
+        task_list: List[DAGNode] = []
+        for task_object in task_or_task_list:

Review comment:
       It is, yes.




-- 
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]


Reply via email to