How to Build a DAG Execution Stage Calculator in Python

from collections import defaultdict, deque


def get_stages(edges):
    """Return list of stages, where each stage is a list of nodes
    that become ready at the same time in a DAG."""
    graph = defaultdict(list)
    in_degree = defaultdict(int)
    nodes = set()

    for src, dst in edges:
        graph[src].append(dst)
        in_degree[dst] += 1
        if src not in in_degree:
            in_degree[src] = 0
        nodes.add(src)
        nodes.add(dst)

    queue = deque([node for node in nodes if in_degree[node] == 0])
    stages = []

    while queue:
        stage = list(queue)
        stages.append(sorted(stage))
        for _ in range(len(queue)):
            node = queue.popleft()
            for neighbor in graph[node]:
                in_degree[neighbor] -= 1
                if in_degree[neighbor] == 0:
                    queue.append(neighbor)

    return stages


if __name__ == "__main__":
    dag_edges = [
        ("parse", "validate"),
        ("validate", "transform"),
        ("transform", "load"),
        ("validate", "aggregate"),
        ("aggregate", "load"),
        ("fetch", "parse"),
    ]
    print(get_stages(dag_edges))