from dataclasses import dataclass
from typing import Dict
from .graph import Graph
from .task import Task
@dataclass(frozen=True)
class TaskGraph:
tasks: Dict[str, Task]
graph: Graph
def __post_init__(self):
assert set(self.tasks) == self.graph.nodes
def for_each_task(self, f, *args, **kwargs):
for task_label in self.graph.visit_postorder():
task = self.tasks[task_label]
f(task, self, *args, **kwargs)
def __getitem__(self, label):
"Get a task by label"
return self.tasks[label]
def __contains__(self, label):
return label in self.tasks
def __iter__(self):
"Iterate over tasks in undefined order"
return iter(self.tasks.values())
def to_json(self):
"Return a JSON-able object representing the task graph, as documented"
named_links_dict = self.graph.named_links_dict()
tasks = {}
for key in self.graph.visit_postorder():
tasks[key] = self.tasks[key].to_json()
tasks[key]["dependencies"] = named_links_dict.get(key, {})
return tasks
@classmethod
def from_json(cls, tasks_dict):
tasks = {}
edges = set()
for key, value in tasks_dict.items():
tasks[key] = Task.from_json(value)
if "task_id" in value:
tasks[key].task_id = value["task_id"]
for depname, dep in value["dependencies"].items():
edges.add((key, dep, depname))
task_graph = cls(tasks, Graph(set(tasks), edges)) return tasks, task_graph