def dag(self) -> Tuple[Dict, Dict]:
"""Construct the DAG of this pipeline based on the its operations and their downstream."""
from polyflow import dags
operations = self.operations.all().prefetch_related('downstream_operations')
def get_downstream(op):
return op.downstream_operations.values_list('id', flat=True)
return dags.get_dag(operations, get_downstream)
def test_get_dag(self):
operations = [OperationFactory() for _ in range(4)]
operations[0].upstream_operations.set(operations[2:])
operations[1].upstream_operations.set(operations[2:])
operation_by_ids = {op.id: op for op in operations}
def get_downstream(op):
return op.downstream_operations.values_list('id', flat=True)
assert dags.get_dag(nodes=operations,
downstream_fn=get_downstream) == ({
operations[0].id:
set(),
operations[1].id:
set(),
operations[2].id:
{operations[0].id, operations[1].id},
operations[3].id:
{operations[0].id, operations[1].id},
}, operation_by_ids)
# Add operations outside the dag
operation1 = OperationFactory()
operation1.downstream_operations.set(
[operations[1], operations[2], operations[3]])
operation2 = OperationFactory()
operation2.upstream_operations.set([operations[0], operations[2]])
assert dags.get_dag(
nodes=operations, downstream_fn=get_downstream) == ({
operations[0].id: {
operation2.id,
},
operations[1].id:
set(),
operations[2].id:
{operations[0].id, operations[1].id, operation2.id},
operations[3].id: {operations[0].id, operations[1].id},
}, operation_by_ids)