Simplify

jedcunningham · ashb · commit cfb7baafa2e9 · 2022-09-07T09:17:57.000+01:00
diff --git a/airflow/models/taskmixin.py b/airflow/models/taskmixin.py
@@ -28,6 +28,7 @@
     from logging import Logger
 
     from airflow.models.dag import DAG
+    from airflow.models.operator import Operator
     from airflow.utils.edgemodifier import EdgeModifier
     from airflow.utils.task_group import TaskGroup
 
@@ -236,14 +237,14 @@ def set_upstream(
         self._set_relatives(task_or_task_list, upstream=True, edge_modifier=edge_modifier)
 
     @property
-    def downstream_list(self) -> Iterable["DAGNode"]:
+    def downstream_list(self) -> Iterable["Operator"]:
         """List of nodes directly downstream"""
         if not self.dag:
             raise AirflowException(f'Operator {self} has not been assigned to a DAG yet')
         return [self.dag.get_task(tid) for tid in self.downstream_task_ids]
 
     @property
-    def upstream_list(self) -> Iterable["DAGNode"]:
+    def upstream_list(self) -> Iterable["Operator"]:
         """List of nodes directly upstream"""
         if not self.dag:
             raise AirflowException(f'Operator {self} has not been assigned to a DAG yet')
diff --git a/airflow/www/views.py b/airflow/www/views.py
@@ -515,7 +515,7 @@ def get_value_from_path(key_path, content):
     return elem
 
 
-def dag_edges(dag):
+def dag_edges(dag: DAG):
     """
     Create the list of edges needed to construct the Graph view.
 
@@ -596,21 +596,17 @@ def collect_edges(task_group):
     # Collect all the edges between individual tasks
     edges = set()
 
-    def get_downstream(task):
-        tasks_to_trace = {(task, frozenset(task.downstream_list))}
-        while tasks_to_trace:
-            tasks_to_trace_next: Set[tuple] = set()
-            for task, children in tasks_to_trace:
-                for child in children:
-                    edge = (task.task_id, child.task_id)
-                    if edge in edges:
-                        continue
-                    edges.add(edge)
-                    tasks_to_trace_next.add((child, frozenset(child.downstream_list)))
-            tasks_to_trace = tasks_to_trace_next
-
-    for root in dag.roots:
-        get_downstream(root)
+    tasks_to_trace: List[Operator] = dag.roots
+    while tasks_to_trace:
+        tasks_to_trace_next: List[Operator] = []
+        for task in tasks_to_trace:
+            for child in task.downstream_list:
+                edge = (task.task_id, child.task_id)
+                if edge in edges:
+                    continue
+                edges.add(edge)
+                tasks_to_trace_next.append(child)
+        tasks_to_trace = tasks_to_trace_next
 
     result = []
     # Build result dicts with the two ends of the edge, plus any extra metadata
