flowrep 0.post0.dev1__py3-none-any.whl

flowrep/workflow.py

@@ -0,0 +1,1256 @@
+import ast
+import builtins
+import copy
+import dataclasses
+import inspect
+import textwrap
+from collections import deque
+from collections.abc import Callable, Iterable
+from functools import cached_property, update_wrapper
+from typing import Any, Generic, TypeVar, cast, get_args, get_origin
+
+import networkx as nx
+from networkx.algorithms.dag import topological_sort
+from semantikon.converter import (
+    get_annotated_type_hints,
+    get_return_expressions,
+    get_return_labels,
+    meta_to_dict,
+    parse_input_args,
+    parse_output_args,
+)
+from semantikon.datastructure import (
+    MISSING,
+    CoreMetadata,
+    Edges,
+    Function,
+    Input,
+    Inputs,
+    Missing,
+    Nodes,
+    Output,
+    Outputs,
+    PortType,
+    TypeMetadata,
+    Workflow,
+)
+
+F = TypeVar("F", bound=Callable[..., object])
+
+
+class FunctionWithWorkflow(Generic[F]):
+    def __init__(self, func: F, workflow: dict[str, object], run) -> None:
+        self.func = func
+        self._semantikon_workflow: dict[str, object] = workflow
+        self.run = run
+        update_wrapper(self, func)  # Copies __name__, __doc__, etc.
+
+    def __call__(self, *args, **kwargs):
+        return self.func(*args, **kwargs)
+
+    def __getattr__(self, item):
+        return getattr(self.func, item)
+
+
+def separate_types(
+    data: dict[str, Any], class_dict: dict[str, type] | None = None
+) -> tuple[dict[str, Any], dict[str, type]]:
+    """
+    Separate types from the data dictionary and store them in a class dictionary.
+    The types inside the data dictionary will be replaced by their name (which
+    would for example make it easier to hash it).
+
+    Args:
+        data (dict[str, Any]): The data dictionary containing nodes and types.
+        class_dict (dict[str, type], optional): A dictionary to store types. It
+            is mainly used due to the recursivity of this function. Defaults to
+            None.
+
+    Returns:
+        tuple: A tuple containing the modified data dictionary and the
+            class dictionary.
+    """
+    data = copy.deepcopy(data)
+    if class_dict is None:
+        class_dict = {}
+    if "nodes" in data:
+        for key, node in data["nodes"].items():
+            child_node, child_class_dict = separate_types(node, class_dict)
+            class_dict.update(child_class_dict)
+            data["nodes"][key] = child_node
+    for io_ in ["inputs", "outputs"]:
+        for key, content in data[io_].items():
+            if "dtype" in content and isinstance(content["dtype"], type):
+                class_dict[content["dtype"].__name__] = content["dtype"]
+                data[io_][key]["dtype"] = content["dtype"].__name__
+    return data, class_dict
+
+
+def separate_functions(
+    data: dict[str, Any], function_dict: dict[str, Callable] | None = None
+) -> tuple[dict[str, Any], dict[str, Callable]]:
+    """
+    Separate functions from the data dictionary and store them in a function
+    dictionary. The functions inside the data dictionary will be replaced by
+    their name (which would for example make it easier to hash it)
+
+    Args:
+        data (dict[str, Any]): The data dictionary containing nodes and
+            functions.
+        function_dict (dict[str, Callable], optional): A dictionary to store
+            functions. It is mainly used due to the recursivity of this
+            function. Defaults to None.
+
+    Returns:
+        tuple: A tuple containing the modified data dictionary and the
+            function dictionary.
+    """
+    data = copy.deepcopy(data)
+    if function_dict is None:
+        function_dict = {}
+    if "nodes" in data:
+        for key, node in data["nodes"].items():
+            child_node, child_function_dict = separate_functions(node, function_dict)
+            function_dict.update(child_function_dict)
+            data["nodes"][key] = child_node
+    elif "function" in data and not isinstance(data["function"], str):
+        fnc_object = data["function"]
+        as_string = fnc_object.__module__ + "." + fnc_object.__qualname__
+        function_dict[as_string] = fnc_object
+        data["function"] = as_string
+    if "test" in data and not isinstance(data["test"]["function"], str):
+        fnc_object = data["test"]["function"]
+        as_string = fnc_object.__module__ + fnc_object.__qualname__
+        function_dict[as_string] = fnc_object
+        data["test"]["function"] = as_string
+    return data, function_dict
+
+
+class FunctionDictFlowAnalyzer:
+    def __init__(self, ast_dict, scope):
+        self.graph = nx.DiGraph()
+        self.scope = scope  # mapping from function names to objects
+        self.function_defs = {}
+        self.ast_dict = ast_dict
+        self._call_counter = {}
+        self._control_flow_list = []
+        self._parallel_var = {}
+
+    def analyze(self) -> tuple[nx.DiGraph, dict[str, Any]]:
+        for arg in self.ast_dict.get("args", {}).get("args", []):
+            if arg["_type"] == "arg":
+                self._add_output_edge("input", arg["arg"])
+        return_was_called = False
+        for node in self.ast_dict.get("body", []):
+            assert not return_was_called
+            self._visit_node(node)
+            if node["_type"] == "Return":
+                return_was_called = True
+        return self.graph, self.function_defs
+
+    def _visit_node(self, node, control_flow: str | None = None):
+        if node["_type"] == "Assign":
+            self._handle_assign(node, control_flow=control_flow)
+        elif node["_type"] == "Expr":
+            self._handle_expr(node, control_flow=control_flow)
+        elif node["_type"] == "While":
+            self._handle_while(node, control_flow=control_flow)
+        elif node["_type"] == "For":
+            self._handle_for(node, control_flow=control_flow)
+        elif node["_type"] == "If":
+            self._handle_if(node, control_flow=control_flow)
+        elif node["_type"] == "Return":
+            self._handle_return(node, control_flow=control_flow)
+        else:
+            raise NotImplementedError(f"Node type {node['_type']} not implemented")
+
+    def _handle_return(self, node, control_flow: str | None = None):
+        if not node["value"]:
+            return
+        if node["value"]["_type"] == "Tuple":
+            for idx, elt in enumerate(node["value"]["elts"]):
+                if elt["_type"] != "Name":
+                    raise NotImplementedError("Only variable returns supported")
+                self._add_input_edge(elt, "output", input_index=idx)
+        elif node["value"]["_type"] == "Name":
+            self._add_input_edge(node["value"], "output")
+
+    def _handle_if(self, node, control_flow: str | None = None):
+        assert node["test"]["_type"] == "Call"
+        control_flow = self._convert_control_flow(control_flow, tag="If")
+        self._parse_function_call(node["test"], control_flow=f"{control_flow}-test")
+        for n in node["body"]:
+            self._visit_node(n, control_flow=f"{control_flow}-body")
+        for n in node.get("orelse", []):
+            cf_else = "/".join(
+                control_flow.split("/")[:-1]
+                + [control_flow.split("/")[-1].replace("If", "Else") + "-body"]
+            )
+            self._visit_node(n, control_flow=cf_else)
+            self._reconnect_parallel(cf_else, f"{control_flow}-body")
+            self._register_parallel_variables(cf_else, f"{control_flow}-body")
+
+    def _reconnect_parallel(self, control_flow: str, ref_control_flow: str):
+        all_edges = list(self.graph.edges.data())
+        for edge in all_edges:
+            if (
+                "control_flow" not in edge[2]
+                or edge[2]["control_flow"] != control_flow
+                or edge[2]["type"] == "output"
+            ):
+                continue
+            var, ind = "_".join(edge[0].split("_")[:-1]), int(edge[0].split("_")[-1])
+            while True:
+                if any(
+                    [
+                        e[2].get("control_flow") == ref_control_flow
+                        for e in self.graph.in_edges(f"{var}_{ind}", data=True)
+                    ]
+                ):
+                    ind -= 1
+                break
+            if f"{var}_{ind}" != edge[0]:
+                self.graph.add_edge(f"{var}_{ind}", edge[1], **edge[2])
+                self.graph.remove_edge(edge[0], edge[1])
+
+    def _register_parallel_variables(self, control_flow: str, ref_control_flow: str):
+        data: dict[str, dict] = {control_flow: {}, ref_control_flow: {}}
+        for edge in self.graph.edges.data():
+            if (
+                edge[2].get("control_flow", "") in [control_flow, ref_control_flow]
+                and edge[2]["type"] == "output"
+            ):
+                data[edge[2]["control_flow"]][edge[1].rsplit("_", 1)[0]] = edge[1]
+        for key in set(data[control_flow].keys()).intersection(
+            data[ref_control_flow].keys()
+        ):
+            values = sorted([data[control_flow][key], data[ref_control_flow][key]])
+            self._parallel_var[values[-1]] = [values[0]]
+            if values[0] in self._parallel_var:
+                self._parallel_var[values[-1]].extend(self._parallel_var.pop(values[0]))
+
+    def _handle_while(self, node, control_flow: str | None = None):
+        assert node["test"]["_type"] == "Call"
+        control_flow = self._convert_control_flow(control_flow, tag="While")
+        self._parse_function_call(node["test"], control_flow=f"{control_flow}-test")
+        for n in node["body"]:
+            self._visit_node(n, control_flow=f"{control_flow}-body")
+
+    def _handle_for(self, node, control_flow: str | None = None):
+        assert node["iter"]["_type"] == "Call"
+        control_flow = self._convert_control_flow(control_flow, tag="For")
+
+        unique_func_name = self._parse_function_call(
+            node["iter"], control_flow=f"{control_flow}-iter"
+        )
+        self._parse_outputs(
+            [node["target"]], unique_func_name, control_flow=control_flow
+        )
+        for n in node["body"]:
+            self._visit_node(n, control_flow=f"{control_flow}-body")
+
+    def _handle_expr(self, node, control_flow: str | None = None) -> str:
+        value = node["value"]
+        return self._parse_function_call(value, control_flow=control_flow)
+
+    def _parse_function_call(self, value, control_flow: str | None = None) -> str:
+        if value["_type"] != "Call":
+            raise NotImplementedError(
+                f"Only function calls allowed on RHS: {value['_type']}"
+            )
+
+        func_node = value["func"]
+        if func_node["_type"] != "Name":
+            raise NotImplementedError("Only simple functions allowed")
+
+        func_name = func_node["id"]
+        unique_func_name = self._get_unique_func_name(func_name)
+
+        if func_name not in self.scope:
+            raise ValueError(f"Function {func_name} not found in scope")
+
+        self.function_defs[unique_func_name] = {"function": self.scope[func_name]}
+        if control_flow is not None:
+            self.function_defs[unique_func_name]["control_flow"] = control_flow
+
+        # Parse inputs (positional + keyword)
+        for i, arg in enumerate(value.get("args", [])):
+            self._add_input_edge(
+                arg, unique_func_name, input_index=i, control_flow=control_flow
+            )
+        for kw in value.get("keywords", []):
+            self._add_input_edge(
+                kw["value"],
+                unique_func_name,
+                input_name=kw["arg"],
+                control_flow=control_flow,
+            )
+        return unique_func_name
+
+    def _handle_assign(self, node, control_flow: str | None = None):
+        unique_func_name = self._handle_expr(node, control_flow=control_flow)
+        self._parse_outputs(
+            node["targets"], unique_func_name, control_flow=control_flow
+        )
+
+    def _parse_outputs(
+        self, targets, unique_func_name, control_flow: str | None = None
+    ):
+        if len(targets) == 1 and targets[0]["_type"] == "Tuple":
+            for idx, elt in enumerate(targets[0]["elts"]):
+                self._add_output_edge(
+                    unique_func_name,
+                    elt["id"],
+                    output_index=idx,
+                    control_flow=control_flow,
+                )
+        else:
+            for target in targets:
+                self._add_output_edge(
+                    unique_func_name, target["id"], control_flow=control_flow
+                )
+
+    def _get_max_index(self, variable: str) -> int:
+        index = 0
+        while True:
+            if len(self.graph.in_edges(f"{variable}_{index}")) > 0:
+                index += 1
+                continue
+            break
+        return index
+
+    def _get_var_index(self, variable: str, output: bool = False) -> int:
+        index = self._get_max_index(variable)
+        if index == 0 and not output:
+            raise KeyError(
+                f"Variable {variable} not found in graph. "
+                "This usually means that the variable was never defined."
+            )
+        if output:
+            return index
+        else:
+            return index - 1
+
+    def _add_output_edge(
+        self, source: str, target: str, control_flow: str | None = None, **kwargs
+    ):
+        """
+        Add an output edge from the source to the target variable.
+
+        Args:
+            source (str): The source node (function name).
+            target (str): The target variable name.
+            control_flow (str | None): The control flow tag, if any.
+            **kwargs: Additional keyword arguments to pass to the edge.
+
+        In the case of the following line:
+
+        >>> y = f(x)
+
+        This function will add an edge from the function `f` to the variable `y`.
+        """
+        versioned = f"{target}_{self._get_var_index(target, output=True)}"
+        if control_flow is not None:
+            kwargs["control_flow"] = control_flow
+        self.graph.add_edge(source, versioned, type="output", **kwargs)
+
+    def _add_input_edge(
+        self, source: dict, target: str, control_flow: str | None = None, **kwargs
+    ):
+        """
+        Add an input edge from the source variable to the target node.
+
+        Args:
+            source (dict): The source variable node.
+            target (str): The target node (function name).
+            control_flow (str | None): The control flow tag, if any.
+            **kwargs: Additional keyword arguments to pass to the edge.
+
+        In the case of the following line:
+
+        >>> y = f(x)
+
+        This function will add an edge from the variable `x` to the function `f`.
+        """
+        if source["_type"] != "Name":
+            raise NotImplementedError(f"Only variable inputs supported, got: {source}")
+        var_name = source["id"]
+        if control_flow is not None:
+            kwargs["control_flow"] = control_flow
+        versioned = f"{var_name}_{self._get_var_index(var_name)}"
+        self.graph.add_edge(versioned, target, type="input", **kwargs)
+        if versioned in self._parallel_var:
+            for key in self._parallel_var.pop(versioned):
+                self.graph.add_edge(key, target, type="input", **kwargs)
+
+    def _get_unique_func_name(self, base_name):
+        i = self._call_counter.get(base_name, 0)
+        self._call_counter[base_name] = i + 1
+        return f"{base_name}_{i}"
+
+    def _convert_control_flow(self, control_flow: str | None, tag: str) -> str:
+        control_flow = "" if control_flow is None else f"{control_flow.split('-')[0]}/"
+        counter = 0
+        while True:
+            if f"{control_flow}{tag}_{counter}" not in self._control_flow_list:
+                self._control_flow_list.append(f"{control_flow}{tag}_{counter}")
+                break
+            counter += 1
+        return f"{control_flow}{tag}_{counter}"
+
+
+def _get_variables_from_subgraph(graph: nx.DiGraph, io_: str) -> set[str]:
+    """
+    Get variables from a subgraph based on the type of I/O and control flow.
+
+    Args:
+        graph (nx.DiGraph): The directed graph representing the function.
+        io_ (str): The type of I/O to filter by ("input" or "output").
+        control_flow (list, str): A list of control flow types to filter by.
+
+    Returns:
+        set[str]: A set of variable names that match the specified I/O type and
+            control flow.
+    """
+    assert io_ in ["input", "output"], "io_ must be 'input' or 'output'"
+    if io_ == "input":
+        edge_ind = 0
+    elif io_ == "output":
+        edge_ind = 1
+    return set(
+        [edge[edge_ind] for edge in graph.edges.data() if edge[2]["type"] == io_]
+    )
+
+
+def _get_parent_graph(graph: nx.DiGraph, control_flow: str) -> nx.DiGraph:
+    """
+    Get parent body of the indented body
+
+    Args:
+        graph (nx.DiGraph): Full graph to look for the parent graph from
+        control_flow (str): Control flow whose parent graph is to look for
+
+    Returns:
+        (nx.DiGraph): Parent graph
+    """
+    return nx.DiGraph(
+        [
+            edge
+            for edge in graph.edges.data()
+            if not _get_control_flow(edge[2]).startswith(control_flow)
+        ]
+    )
+
+
+def _detect_io_variables_from_control_flow(
+    graph: nx.DiGraph, subgraph: nx.DiGraph
+) -> dict[str, list]:
+    """
+    Detect input and output variables from a graph based on control flow.
+
+    Args:
+        graph (nx.DiGraph): The directed graph representing the function.
+
+    Returns:
+        dict[str, set]: A dictionary with keys "input" and "output", each
+            containing a set
+
+    Take a look at the unit tests for examples of how to use this function.
+    """
+    sg_body = nx.DiGraph(
+        [
+            edge
+            for edge in subgraph.edges.data()
+            if edge[0] != "input" and edge[1] != "output"
+        ]
+    )
+    cf = sorted(
+        [
+            _get_control_flow(edge[2])
+            for edge in sg_body.edges.data()
+            if "control_flow" in edge[2]
+        ]
+    )
+    if len(cf) == 0:
+        return {"inputs": [], "outputs": []}
+    assert all([cf[ii + 1].startswith(cf[ii]) for ii in range(len(cf) - 1)])
+    parent_graph = _get_parent_graph(graph, cf[0])
+    var_inp_1 = _get_variables_from_subgraph(graph=sg_body, io_="input")
+    var_inp_2 = _get_variables_from_subgraph(graph=parent_graph, io_="output")
+    var_out_1 = _get_variables_from_subgraph(graph=parent_graph, io_="input")
+    var_out_2 = _get_variables_from_subgraph(graph=sg_body, io_="output")
+    return {
+        "inputs": list(var_inp_1.intersection(var_inp_2)),
+        "outputs": list(var_out_1.intersection(var_out_2)),
+    }
+
+
+def _extract_control_flows(graph: nx.DiGraph) -> list[str]:
+    return list(set([_get_control_flow(edge[2]) for edge in graph.edges.data()]))
+
+
+def _split_graphs_into_subgraphs(graph: nx.DiGraph) -> dict[str, nx.DiGraph]:
+    return {
+        control_flow: nx.DiGraph(
+            [
+                edge
+                for edge in graph.edges.data()
+                if _get_control_flow(edge[2]) == control_flow
+            ]
+        )
+        for control_flow in _extract_control_flows(graph)
+    }
+
+
+def _get_subgraphs(graph: nx.DiGraph, cf_graph: nx.DiGraph) -> dict[str, nx.DiGraph]:
+    """
+    Separate a flat graph into subgraphs nested by control flows
+
+    Args:
+        graph (nx.DiGraph): Flat workflow graph
+        cf_graph (nx.DiGraph): Control flow graph (cf. _get_control_flow_graph)
+
+    Returns:
+        dict[str, nx.DiGraph]: Subgraphs
+    """
+    subgraphs = _split_graphs_into_subgraphs(graph)
+    for key in list(topological_sort(cf_graph))[::-1]:
+        subgraph = subgraphs[key]
+        node_name = "injected_" + key.replace("/", "_")
+        io_ = _detect_io_variables_from_control_flow(graph, subgraph)
+        for parent_graph_name in cf_graph.predecessors(key):
+            parent_graph = subgraphs[parent_graph_name]
+            for inp in io_["inputs"]:
+                parent_graph.add_edge(
+                    inp, node_name, type="input", input_name=_remove_index(inp)
+                )
+            for out in io_["outputs"]:
+                parent_graph.add_edge(
+                    node_name, out, type="output", output_name=_remove_index(out)
+                )
+        for inp in io_["inputs"]:
+            subgraph.add_edge("input", inp, type="output")
+        for out in io_["outputs"]:
+            subgraph.add_edge(out, "output", type="input")
+    return subgraphs
+
+
+def _extract_functions_from_graph(graph: nx.DiGraph) -> set:
+    function_names = []
+    for edge in graph.edges.data():
+        if edge[2]["type"] == "output" and edge[0] != "input":
+            function_names.append(edge[0])
+        elif edge[2]["type"] == "input" and edge[1] != "output":
+            function_names.append(edge[1])
+    return set(function_names)
+
+
+def _get_control_flow_graph(control_flows: list[str]) -> nx.DiGraph:
+    """
+    Create a graph based on the control flows. The indentation level
+    corresponds to the graph level. The higher level body is the parent node
+    of the lower body.
+
+
+    Args:
+        control_flows (list[str]): All control flows present in a workflow
+
+    Returns:
+        nx.DiGraph: Control flow graph
+    """
+    cf_list = []
+    for cf in control_flows:
+        if cf == "":
+            continue
+        if "/" in cf:
+            cf_list.append(["/".join(cf.split("/")[:-1]), cf])
+        else:
+            cf_list.append(["", cf])
+    graph = nx.DiGraph(cf_list)
+    if len(graph) == 0:
+        graph.add_node("")
+    return graph
+
+
+def _function_to_ast_dict(node):
+    if isinstance(node, ast.AST):
+        result = {"_type": type(node).__name__}
+        for field, value in ast.iter_fields(node):
+            result[field] = _function_to_ast_dict(value)
+        return result
+    elif isinstance(node, list):
+        return [_function_to_ast_dict(item) for item in node]
+    else:
+        return node
+
+
+def get_ast_dict(func: Callable) -> dict:
+    """Get the AST dictionary representation of a function."""
+    source_code = textwrap.dedent(inspect.getsource(func))
+    tree = ast.parse(source_code)
+    return _function_to_ast_dict(tree)
+
+
+def analyze_function(func: Callable) -> tuple[nx.DiGraph, dict[str, Any]]:
+    """Extracts the variable flow graph from a function"""
+    ast_dict = get_ast_dict(func)
+    scope = inspect.getmodule(func).__dict__ | vars(builtins)
+    analyzer = FunctionDictFlowAnalyzer(ast_dict["body"][0], scope)
+    return analyzer.analyze()
+
+
+def _get_node_outputs(func: Callable, counts: int | None = None) -> dict[str, dict]:
+    output_hints = parse_output_args(
+        func, separate_tuple=(counts is None or counts > 1)
+    )
+    output_vars = get_return_expressions(func)
+    if output_vars is None or len(output_vars) == 0:
+        return {}
+    if (counts is not None and counts == 1) or isinstance(output_vars, str):
+        if isinstance(output_vars, str):
+            return {output_vars: cast(dict, output_hints)}
+        else:
+            return {"output": cast(dict, output_hints)}
+    assert isinstance(output_vars, tuple), output_vars
+    assert counts is None or len(output_vars) == counts, output_vars
+    if output_hints == {}:
+        return {key: {} for key in output_vars}
+    else:
+        assert counts is None or len(output_hints) == counts
+        return {key: hint for key, hint in zip(output_vars, output_hints, strict=False)}
+
+
+def _get_output_counts(graph: nx.DiGraph) -> dict[str, int]:
+    """
+    Get the number of outputs for each node in the graph.
+
+    Args:
+        graph (nx.DiGraph): The directed graph representing the function.
+
+    Returns:
+        dict: A dictionary mapping node names to the number of outputs.
+    """
+    f_dict: dict[str, int] = {}
+    for edge in graph.edges.data():
+        if edge[2]["type"] != "output":
+            continue
+        f_dict[edge[0]] = f_dict.get(edge[0], 0) + 1
+    if "input" in f_dict:
+        del f_dict["input"]
+    return f_dict
+
+
+def _get_nodes(
+    data: dict[str, dict],
+    output_counts: dict[str, int],
+    control_flow: None | str = None,
+) -> dict[str, dict]:
+    result = {}
+    for label, function in data.items():
+        func = function["function"]
+        if hasattr(func, "_semantikon_workflow"):
+            if output_counts[label] != len(func._semantikon_workflow["outputs"]):
+                raise ValueError(
+                    f"{label} has {len(func._semantikon_workflow['outputs'])} outputs, "
+                    f"but {output_counts[label]} expected"
+                )
+            data_dict = func._semantikon_workflow.copy()
+            result[label] = data_dict
+            result[label]["label"] = label
+            if hasattr(func, "_semantikon_metadata"):
+                result[label].update(func._semantikon_metadata)
+        else:
+            result[label] = get_node_dict(
+                function=func,
+                inputs=parse_input_args(func),
+                outputs=_get_node_outputs(func, output_counts.get(label, 1)),
+            )
+    return result
+
+
+def _remove_index(s: str) -> str:
+    return "_".join(s.split("_")[:-1])
+
+
+def _get_control_flow(data: dict[str, Any]) -> str:
+    """
+    Get the control flow name
+
+    Args:
+        data (dict[str, Any]): metadata of the edge (which is stored in the
+            third element of each edge of nx.Digraph)
+
+    Returns:
+        (str): Control flow name (e.g. While_0, For_3 etc.)
+    """
+    return data.get("control_flow", "").split("-")[0]
+
+
+def _get_sorted_edges(graph: nx.DiGraph) -> list:
+    """
+    Sort the edges of the graph based on the topological order of the nodes.
+
+    Args:
+        graph (nx.DiGraph): The directed graph representing the function.
+
+    Returns:
+        list: A sorted list of edges in the graph.
+
+    Example:
+
+    >>> graph.add_edges_from([('A', 'B'), ('B', 'D'), ('A', 'C'), ('C', 'D')])
+    >>> sorted_edges = _get_sorted_edges(graph)
+    >>> print(sorted_edges)
+
+    Output:
+
+    >>> [('A', 'B', {}), ('A', 'C', {}), ('B', 'D', {}), ('C', 'D', {})]
+    """
+    topo_order = list(topological_sort(graph))
+    node_order = {node: i for i, node in enumerate(topo_order)}
+    return sorted(graph.edges.data(), key=lambda edge: node_order[edge[0]])
+
+
+def _remove_and_reconnect_nodes(
+    G: nx.DiGraph, nodes_to_remove: list[str]
+) -> nx.DiGraph:
+    for node in set(nodes_to_remove):
+        preds = list(G.predecessors(node))
+        succs = list(G.successors(node))
+        for u in preds:
+            for v in succs:
+                G.add_edge(u, v)
+        G.remove_node(node)
+    return G
+
+
+def _get_edges(graph: nx.DiGraph, nodes: dict[str, dict]) -> list[tuple[str, str]]:
+    io_dict = {
+        key: {
+            "input": list(data["inputs"].keys()),
+            "output": list(data["outputs"].keys()),
+        }
+        for key, data in nodes.items()
+    }
+    edges = []
+    nodes_to_remove = []
+    for edge in graph.edges.data():
+        if edge[0] == "input":
+            edges.append([edge[0] + "s." + _remove_index(edge[1]), edge[1]])
+            nodes_to_remove.append(edge[1])
+        elif edge[1] == "output":
+            edges.append([edge[0], edge[1] + "s." + _remove_index(edge[0])])
+            nodes_to_remove.append(edge[0])
+        elif edge[2]["type"] == "input":
+            if "input_name" in edge[2]:
+                tag = edge[2]["input_name"]
+            elif "input_index" in edge[2]:
+                tag = io_dict[edge[1]]["input"][edge[2]["input_index"]]
+            else:
+                raise ValueError
+            edges.append([edge[0], edge[1] + ".inputs." + tag])
+            nodes_to_remove.append(edge[0])
+        elif edge[2]["type"] == "output":
+            if "output_index" in edge[2]:
+                tag = io_dict[edge[0]]["output"][edge[2]["output_index"]]
+            elif "output_name" in edge[2]:
+                tag = edge[2]["output_name"]
+            else:
+                tag = io_dict[edge[0]]["output"][0]
+            edges.append([edge[0] + ".outputs." + tag, edge[1]])
+            nodes_to_remove.append(edge[1])
+    new_graph = _remove_and_reconnect_nodes(nx.DiGraph(edges), nodes_to_remove)
+    return list(new_graph.edges)
+
+
+def get_node_dict(
+    function: Callable,
+    inputs: dict[str, dict] | None = None,
+    outputs: dict[str, dict] | None = None,
+) -> dict:
+    """
+    Get a dictionary representation of the function node.
+
+    Args:
+        func (Callable): The function to be analyzed.
+        data_format (str): The format of the output. Options are "semantikon" and
+            "ape".
+
+    Returns:
+        (dict) A dictionary representation of the function node.
+    """
+    if inputs is None:
+        inputs = parse_input_args(function)
+    if outputs is None:
+        outputs = _get_node_outputs(function)
+    data = {
+        "inputs": inputs,
+        "outputs": outputs,
+        "function": function,
+        "type": "Function",
+    }
+    if hasattr(function, "_semantikon_metadata"):
+        data.update(function._semantikon_metadata)
+    return data
+
+
+def _to_workflow_dict_entry(
+    inputs: dict[str, dict],
+    outputs: dict[str, dict],
+    nodes: dict[str, dict],
+    edges: list[tuple[str, str]],
+    label: str,
+    **kwargs,
+) -> dict[str, object]:
+    assert all("inputs" in v for v in nodes.values())
+    assert all("outputs" in v for v in nodes.values())
+    assert all(
+        "function" in v or ("nodes" in v and "edges" in v) for v in nodes.values()
+    )
+    return {
+        "inputs": inputs,
+        "outputs": outputs,
+        "nodes": nodes,
+        "edges": edges,
+        "label": label,
+        "type": "Workflow",
+    } | kwargs
+
+
+def _get_test_dict(f_dict: dict[str, dict]) -> dict[str, str]:
+    """
+    dict to translate test and iter nodes into "test" and "iter"
+
+    Args:
+        f_dict (dict[str, dict]): Function dictionary
+
+    Returns:
+        dict[str, str]: Translation of node name to "test" or "iter"
+    """
+    return {
+        key: tag
+        for key, value in f_dict.items()
+        for tag in ["test", "iter"]
+        if value.get("control_flow", "").endswith(tag)
+    }
+
+
+def _nest_nodes(
+    graph: nx.DiGraph, nodes: dict[str, dict], f_dict: dict[str, dict]
+) -> tuple[dict[str, dict], list[tuple[str, str]]]:
+    """
+    Nest workflow nodes
+
+    Args:
+        graph (nx.DiGraph): The directed graph representing the function.
+        nodes (dict[str, dict]): The dictionary of nodes.
+        f_dict (dict[str, dict]): The dictionary of functions.
+
+    Returns:
+        dict: A dictionary containing the nested nodes, edges, and label.
+    """
+    test_dict = _get_test_dict(f_dict=f_dict)
+    cf_graph = _get_control_flow_graph(_extract_control_flows(graph))
+    subgraphs = _get_subgraphs(graph, cf_graph)
+    injected_nodes: dict[str, Any] = {}
+    for cf_key in list(topological_sort(cf_graph))[::-1]:
+        subgraph = nx.relabel_nodes(subgraphs[cf_key], test_dict)
+        new_key = "injected_" + cf_key.replace("/", "_") if len(cf_key) > 0 else cf_key
+        current_nodes = {}
+        for key in _extract_functions_from_graph(subgraphs[cf_key]):
+            if key in test_dict:
+                current_nodes[test_dict[key]] = nodes[key]
+            elif key in nodes:
+                current_nodes[key] = nodes[key]
+            else:
+                current_nodes[key] = injected_nodes.pop(key)
+        io_ = _detect_io_variables_from_control_flow(graph, subgraph)
+        injected_nodes[new_key] = {
+            "nodes": current_nodes,
+            "edges": _get_edges(subgraph, current_nodes),
+            "label": new_key,
+            "inputs": {_remove_index(key): {} for key in io_["inputs"]},
+            "outputs": {_remove_index(key): {} for key in io_["outputs"]},
+        }
+        for tag in ["test", "iter"]:
+            if tag in injected_nodes[new_key]["nodes"]:
+                injected_nodes[new_key][tag] = injected_nodes[new_key]["nodes"].pop(tag)
+    return injected_nodes[""]["nodes"], injected_nodes[""]["edges"]
+
+
+def get_workflow_dict(func: Callable) -> dict[str, object]:
+    """
+    Get a dictionary representation of the workflow for a given function.
+
+    Args:
+        func (Callable): The function to be analyzed.
+
+    Returns:
+        dict: A dictionary representation of the workflow, including inputs,
+            outputs, nodes, edges, and label.
+    """
+    graph, f_dict = analyze_function(func)
+    nodes = _get_nodes(f_dict, _get_output_counts(graph))
+    nested_nodes, edges = _nest_nodes(graph, nodes, f_dict)
+    return _to_workflow_dict_entry(
+        inputs=parse_input_args(func),
+        outputs=_get_node_outputs(func),
+        nodes=nested_nodes,
+        edges=edges,
+        label=func.__name__,
+    )
+
+
+def _get_missing_edges(edge_list: list[tuple[str, str]]) -> list[tuple[str, str]]:
+    """
+    Insert processes into the data edges. Take the following workflow:
+
+    >>> y = f(x=x)
+    >>> z = g(y=y)
+
+    The data flow is
+
+    - f.inputs.x -> f.outputs.y
+    - f.outputs.y -> g.inputs.y
+    - g.inputs.y -> g.outputs.z
+
+    `_get_missing_edges` adds the processes:
+
+    - f.inputs.x -> f
+    - f -> f.outputs.y
+    - f.outputs.y -> g.inputs.y
+    - g.inputs.y -> g
+    - g -> g.outputs.z
+    """
+    extra_edges = []
+    for edge in edge_list:
+        for tag in edge:
+            if len(tag.split(".")) < 3:
+                continue
+            if tag.split(".")[1] == "inputs":
+                new_edge = (tag, tag.split(".")[0])
+            elif tag.split(".")[1] == "outputs":
+                new_edge = (tag.split(".")[0], tag)
+            if new_edge not in extra_edges:
+                extra_edges.append(new_edge)
+    return extra_edges
+
+
+class _Workflow:
+    def __init__(self, workflow_dict: dict[str, Any]):
+        self._workflow = workflow_dict
+
+    @cached_property
+    def _all_edges(self) -> list[tuple[str, str]]:
+        edges = cast(dict[str, list], self._workflow)["edges"]
+        return edges + _get_missing_edges(edges)
+
+    @cached_property
+    def _graph(self) -> nx.DiGraph:
+        return nx.DiGraph(self._all_edges)
+
+    @cached_property
+    def _execution_list(self) -> list[list[str]]:
+        return find_parallel_execution_levels(self._graph)
+
+    def _sanitize_input(self, *args, **kwargs) -> dict[str, Any]:
+        keys = list(self._workflow["inputs"].keys())
+        for ii, arg in enumerate(args):
+            if keys[ii] in kwargs:
+                raise TypeError(
+                    f"{self._workflow['label']}() got multiple values for"
+                    " argument '{keys[ii]}'"
+                )
+            kwargs[keys[ii]] = arg
+        return kwargs
+
+    def _set_inputs(self, *args, **kwargs):
+        kwargs = self._sanitize_input(*args, **kwargs)
+        for key, value in kwargs.items():
+            if key not in self._workflow["inputs"]:
+                raise TypeError(f"Unexpected keyword argument '{key}'")
+            self._workflow["inputs"][key]["value"] = value
+
+    def _get_value_from_data(self, node: dict[str, Any]) -> Any:
+        if "value" not in node:
+            node["value"] = node["default"]
+        return node["value"]
+
+    def _get_value_from_global(self, path: str) -> Any:
+        io, var = path.split(".")
+        return self._get_value_from_data(self._workflow[io][var])
+
+    def _get_value_from_node(self, path: str) -> Any:
+        node, io, var = path.split(".")
+        return self._get_value_from_data(self._workflow["nodes"][node][io][var])
+
+    def _set_value_from_global(self, path, value):
+        io, var = path.split(".")
+        self._workflow[io][var]["value"] = value
+
+    def _set_value_from_node(self, path, value):
+        node, io, var = path.split(".")
+        try:
+            self._workflow["nodes"][node][io][var]["value"] = value
+        except KeyError:
+            raise KeyError(f"{path} not found in {node}") from None
+
+    def _execute_node(self, function: str) -> Any:
+        node = self._workflow["nodes"][function]
+        input_data = {}
+        try:
+            for key, content in node["inputs"].items():
+                if "value" not in content:
+                    content["value"] = content["default"]
+                input_data[key] = content["value"]
+        except KeyError:
+            raise KeyError(f"value not defined for {function}") from None
+        if "function" not in node:
+            workflow = _Workflow(node)
+            outputs = [
+                d["value"] for d in workflow.run(**input_data)["outputs"].values()
+            ]
+            if len(outputs) == 1:
+                outputs = outputs[0]
+        else:
+            outputs = node["function"](**input_data)
+        return outputs
+
+    def _set_value(self, tag, value):
+        if len(tag.split(".")) == 2 and tag.split(".")[0] in ("inputs", "outputs"):
+            self._set_value_from_global(tag, value)
+        elif len(tag.split(".")) == 3 and tag.split(".")[1] in ("inputs", "outputs"):
+            self._set_value_from_node(tag, value)
+        elif "." in tag:
+            raise ValueError(f"{tag} not recognized")
+
+    def _get_value(self, tag: str):
+        if len(tag.split(".")) == 2 and tag.split(".")[0] in ("inputs", "outputs"):
+            return self._get_value_from_global(tag)
+        elif len(tag.split(".")) == 3 and tag.split(".")[1] in ("inputs", "outputs"):
+            return self._get_value_from_node(tag)
+        elif "." not in tag:
+            return self._execute_node(tag)
+        else:
+            raise ValueError(f"{tag} not recognized")
+
+    def run(self, *args, **kwargs) -> dict[str, Any]:
+        self._set_inputs(*args, **kwargs)
+        for current_list in self._execution_list:
+            for item in current_list:
+                values = self._get_value(item)
+                nodes = self._graph.edges(item)
+                if "." not in item and len(nodes) > 1:
+                    for value, node in zip(values, nodes, strict=False):
+                        self._set_value(node[1], value)
+                else:
+                    for node in nodes:
+                        self._set_value(node[1], values)
+        return self._workflow
+
+
+def find_parallel_execution_levels(G: nx.DiGraph) -> list[list[str]]:
+    """
+    Find levels of parallel execution in a directed acyclic graph (DAG).
+
+    Args:
+        G (nx.DiGraph): The directed graph representing the function.
+
+    Returns:
+        list[list[str]]: A list of lists, where each inner list contains nodes
+            that can be executed in parallel.
+
+    Comment:
+        This function only gives you a list of nodes that can be executed in
+        parallel, but does not tell you which processes can be executed in
+        case there is a process that takes longer at a higher level.
+    """
+    in_degree = dict(cast(Iterable[tuple[Any, int]], G.in_degree()))
+    queue = deque([node for node in G.nodes if in_degree[node] == 0])
+    levels = []
+
+    while queue:
+        current_level = list(queue)
+        if "input" not in current_level and "output" not in current_level:
+            levels.append(current_level)
+
+        next_queue: deque = deque()
+        for node in current_level:
+            for neighbor in G.successors(node):
+                in_degree[neighbor] -= 1
+                if in_degree[neighbor] == 0:
+                    next_queue.append(neighbor)
+
+        queue = next_queue
+
+    return levels
+
+
+def workflow(func: Callable) -> FunctionWithWorkflow:
+    """
+    Decorator to convert a function into a workflow with metadata.
+
+    Args:
+        func (Callable): The function to be converted into a workflow.
+
+    Returns:
+        FunctionWithWorkflow: A callable object that includes the original function
+
+    Example:
+
+    >>> def operation(x: float, y: float) -> tuple[float, float]:
+    >>>     return x + y, x - y
+    >>>
+    >>>
+    >>> def add(x: float = 2.0, y: float = 1) -> float:
+    >>>     return x + y
+    >>>
+    >>>
+    >>> def multiply(x: float, y: float = 5) -> float:
+    >>>     return x * y
+    >>>
+    >>>
+    >>> @workflow
+    >>> def example_macro(a=10, b=20):
+    >>>     c, d = operation(a, b)
+    >>>     e = add(c, y=d)
+    >>>     f = multiply(e)
+    >>>     return f
+    >>>
+    >>>
+    >>> @workflow
+    >>> def example_workflow(a=10, b=20):
+    >>>     y = example_macro(a, b)
+    >>>     z = add(y, b)
+    >>>     return z
+
+    This example defines a workflow `example_macro`, that includes `operation`,
+    `add`, and `multiply`, which is nested inside another workflow
+    `example_workflow`. Both workflows can be executed using their `run` method,
+    which returns the dictionary representation of the workflow with all the
+    intermediate steps and outputs.
+    """
+    workflow_dict = get_workflow_dict(func)
+    w = _Workflow(workflow_dict)
+    func_with_metadata = FunctionWithWorkflow(func, workflow_dict, w.run)
+    return func_with_metadata
+
+
+def get_ports(
+    func: Callable, separate_return_tuple: bool = True, strict: bool = False
+) -> tuple[Inputs, Outputs]:
+    type_hints = get_annotated_type_hints(func)
+    return_hint = type_hints.pop("return", inspect.Parameter.empty)
+    return_labels = get_return_labels(
+        func, separate_tuple=separate_return_tuple, strict=strict
+    )
+    if get_origin(return_hint) is tuple and separate_return_tuple:
+        output_annotations = {
+            label: meta_to_dict(ann, flatten_metadata=False)
+            for label, ann in zip(return_labels, get_args(return_hint), strict=False)
+        }
+    else:
+        output_annotations = {
+            return_labels[0]: meta_to_dict(return_hint, flatten_metadata=False)
+        }
+    input_annotations = {
+        key: meta_to_dict(
+            type_hints.get(key, value.annotation), value.default, flatten_metadata=False
+        )
+        for key, value in inspect.signature(func).parameters.items()
+    }
+    return (
+        Inputs(**{k: Input(label=k, **v) for k, v in input_annotations.items()}),
+        Outputs(**{k: Output(label=k, **v) for k, v in output_annotations.items()}),
+    )
+
+
+def get_node(func: Callable, label: str | None = None) -> Function | Workflow:
+    metadata_dict = (
+        func._semantikon_metadata if hasattr(func, "_semantikon_metadata") else MISSING
+    )
+    metadata = (
+        metadata_dict
+        if isinstance(metadata_dict, Missing)
+        else CoreMetadata.from_dict(metadata_dict)
+    )
+
+    if hasattr(func, "_semantikon_workflow"):
+        return parse_workflow(func._semantikon_workflow, metadata)
+    else:
+        return parse_function(func, metadata, label=label)
+
+
+def parse_function(
+    func: Callable, metadata: CoreMetadata | Missing, label: str | None = None
+) -> Function:
+    inputs, outputs = get_ports(func)
+    return Function(
+        label=func.__name__ if label is None else label,
+        inputs=inputs,
+        outputs=outputs,
+        function=func,
+        metadata=metadata,
+    )
+
+
+def _port_from_dictionary(
+    io_dictionary: dict[str, object], label: str, port_class: type[PortType]
+) -> PortType:
+    """
+    Take a traditional _semantikon_workflow dictionary's input or output subdictionary
+    and nest the metadata (if any) as a dataclass.
+    """
+    metadata_kwargs = {}
+    for field in dataclasses.fields(TypeMetadata):
+        if field.name in io_dictionary:
+            metadata_kwargs[field.name] = io_dictionary.pop(field.name)
+    if len(metadata_kwargs) > 0:
+        io_dictionary["metadata"] = TypeMetadata.from_dict(metadata_kwargs)
+    io_dictionary["label"] = label
+    return port_class.from_dict(io_dictionary)
+
+
+def _input_from_dictionary(io_dictionary: dict[str, object], label: str) -> Input:
+    return _port_from_dictionary(io_dictionary, label, Input)
+
+
+def _output_from_dictionary(io_dictionary: dict[str, object], label: str) -> Output:
+    return _port_from_dictionary(io_dictionary, label, Output)
+
+
+def parse_workflow(
+    semantikon_workflow: dict[str, Any], metadata: CoreMetadata | Missing = MISSING
+) -> Workflow:
+    label = semantikon_workflow["label"]
+    inputs = Inputs(
+        **{
+            k: _input_from_dictionary(v, label=k)
+            for k, v in semantikon_workflow["inputs"].items()
+        }
+    )
+    outputs = Outputs(
+        **{
+            k: _output_from_dictionary(v, label=k)
+            for k, v in semantikon_workflow["outputs"].items()
+        }
+    )
+    nodes = Nodes(
+        **{
+            k: (
+                get_node(v["function"], label=k)
+                if "function" in v
+                else parse_workflow(v)
+            )
+            for k, v in semantikon_workflow["nodes"].items()
+        }
+    )
+    edges = Edges(**{v: k for k, v in semantikon_workflow["edges"]})
+    return Workflow(
+        label=label,
+        inputs=inputs,
+        outputs=outputs,
+        nodes=nodes,
+        edges=edges,
+        metadata=metadata,
+    )