dora.builder

  1from __future__ import annotations
  2
  3import yaml
  4
  5
  6class DataflowBuilder:
  7    """A dora dataflow."""
  8
  9    def __init__(self, name: str = "dora-dataflow"):
 10        self.name = name
 11        self.nodes = []
 12
 13    def add_node(self, id: str, **kwargs) -> Node:
 14        """Adds a new node to the dataflow."""
 15        node = Node(id, **kwargs)
 16        self.nodes.append(node)
 17        return node
 18
 19    def to_yaml(self, path: str = None) -> str | None:
 20        """Generates the YAML representation of the dataflow."""
 21        dataflow_spec = {"nodes": [node.to_dict() for node in self.nodes]}
 22        if path:
 23            with open(path, "w") as f:
 24                yaml.dump(dataflow_spec, f)
 25            return None
 26        else:
 27            return yaml.dump(dataflow_spec)
 28
 29    def __enter__(self):
 30        return self
 31
 32    def __exit__(self, exc_type, exc_val, exc_tb):
 33        pass
 34
 35
 36class Output:
 37    """Represents an output from a node."""
 38
 39    def __init__(self, node: Node, output_id: str):
 40        self.node = node
 41        self.output_id = output_id
 42
 43    def __str__(self) -> str:
 44        return f"{self.node.id}/{self.output_id}"
 45
 46
 47class Node:
 48    """A node in a dora dataflow."""
 49
 50    def __init__(self, id: str, **kwargs):
 51        self.id = id
 52        self.config = kwargs
 53        self.config["id"] = id
 54        self.operators = []
 55
 56    def path(self, path: str) -> Node:
 57        """Sets the path to the executable or script."""
 58        self.config["path"] = path
 59        return self
 60
 61    def args(self, args: str) -> Node:
 62        """Sets the command-line arguments for the node."""
 63        self.config["args"] = args
 64        return self
 65
 66    def env(self, env: dict) -> Node:
 67        """Sets the environment variables for the node."""
 68        self.config["env"] = env
 69        return self
 70
 71    def build(self, build_command: str) -> Node:
 72        """Sets the build command for the node."""
 73        self.config["build"] = build_command
 74        return self
 75
 76    def git(
 77        self, url: str, branch: str = None, tag: str = None, rev: str = None
 78    ) -> Node:
 79        """Sets the Git repository for the node."""
 80        self.config["git"] = url
 81        if branch:
 82            self.config["branch"] = branch
 83        if tag:
 84            self.config["tag"] = tag
 85        if rev:
 86            self.config["rev"] = rev
 87        return self
 88
 89    def add_operator(self, operator: Operator) -> Node:
 90        """Adds an operator to this node."""
 91        self.operators.append(operator)
 92        return self
 93
 94    def add_output(self, output_id: str) -> Output:
 95        """Adds an output to the node and returns an Output object."""
 96        if "outputs" not in self.config:
 97            self.config["outputs"] = []
 98        if output_id not in self.config["outputs"]:
 99            self.config["outputs"].append(output_id)
100        return Output(self, output_id)
101
102    def add_input(
103        self, input_id: str, source: str | Output, queue_size: int = None
104    ) -> Node:
105        """Adds a user-defined input to the node. Source can be a string or an Output object."""
106        if "inputs" not in self.config:
107            self.config["inputs"] = {}
108
109        if isinstance(source, Output):
110            source_str = str(source)
111            if queue_size is not None:
112                self.config["inputs"][input_id] = {
113                    "source": source_str,
114                    "queue_size": queue_size,
115                }
116            else:
117                self.config["inputs"][input_id] = source_str
118        else:
119            if queue_size is not None:
120                self.config["inputs"][input_id] = {
121                    "source": source,
122                    "queue_size": queue_size,
123                }
124            else:
125                self.config["inputs"][input_id] = source
126        return self
127
128    def to_dict(self) -> dict:
129        """Returns the dictionary representation of the node."""
130        config = self.config.copy()
131        if self.operators:
132            config["operators"] = [op.to_dict() for op in self.operators]
133        return config
134
135
136class Operator:
137    """An operator in a dora dataflow."""
138
139    def __init__(
140        self,
141        id: str,
142        name: str = None,
143        description: str = None,
144        build: str = None,
145        python: str = None,
146        shared_library: str = None,
147        send_stdout_as: str = None,
148    ):
149        self.id = id
150        self.config = {"id": id}
151        if name:
152            self.config["name"] = name
153        if description:
154            self.config["description"] = description
155        if build:
156            self.config["build"] = build
157        if python:
158            self.config["python"] = python
159        if shared_library:
160            self.config["shared-library"] = shared_library
161        if send_stdout_as:
162            self.config["send_stdout_as"] = send_stdout_as
163
164    def to_dict(self) -> dict:
165        """Returns the dictionary representation of the operator."""
166        return self.config