Source code for ewokscore.graph.graph_io

from typing import Dict, Iterator, List, Mapping, Optional, Union
import networkx

from .analysis import start_nodes
from .analysis import end_nodes
from ..node import NodeIdType
from ..node import get_node_label
from .. import missing_data
from ..task import Task




def update_default_inputs(
    graph: networkx.DiGraph, inputs: Optional[List[dict]] = None
) -> None:
    """Input items have the following keys:

    - name: input variable name
    - value: input variable value
    - id (optional): node id
    - label (optional): used when `id` is missing
    - task_identifier (optional): used when `id` is missing
    - all (optional): used when `id`, `label` and `task_identifier` are missing (`True`: all nodes, `False`: start nodes)
    """
    inputs = parse_inputs(graph, inputs)
    keys_to_update = "name", "value"
    for input_item in inputs:
        node_id = input_item.get("id")
        if node_id is None:
            continue
        node_attrs = graph.nodes[node_id]
        default_inputs = node_attrs.get("default_inputs")
        input_item = {k: input_item[k] for k in keys_to_update}
        if default_inputs:
            for existing_input_item in default_inputs:
                if existing_input_item["name"] == input_item["name"]:
                    existing_input_item.update(input_item)
                    break
            else:
                default_inputs.append(input_item)
        else:
            node_attrs["default_inputs"] = [input_item]





def parse_inputs(
    graph: networkx.DiGraph, inputs: Optional[List[dict]] = None
) -> List[dict]:
    """Input items have the following keys:

    - name: input variable name
    - value: input variable value
    - id (optional): node id
    - label (optional): used when `id` is missing
    - task_identifier (optional): used when `id` is missing
    - all (optional): used when `id`, `label` and `task_identifier` are missing (`True`: all nodes, `False`: start nodes)
    """
    if not inputs:
        return list()
    required = {"name", "value"}
    returned = {"id", "name", "value"}
    parsed = list()
    for input_item in list(inputs):
        missing = required - input_item.keys()
        if missing:
            raise ValueError(f"missing keys in one of the graph inputs: {missing}")
        if "id" in input_item:
            parsed.append({k: v for k, v in input_item.items() if k in returned})
            continue

        node_filters = dict()
        for k in ("label", "task_identifier"):
            if k in input_item:
                node_filters[k] = input_item[k]

        if node_filters:
            node_ids = iter_node_ids(graph, **node_filters)
        elif input_item.get("all"):
            node_ids = graph.nodes
        else:
            node_ids = start_nodes(graph)

        for node_id in node_ids:
            input_item = {k: v for k, v in input_item.items() if k in returned}
            input_item["id"] = node_id
            parsed.append(input_item)
    return parsed





def parse_outputs(
    graph: networkx.DiGraph, outputs: Optional[List[dict]] = None
) -> List[dict]:
    """Output items have the following keys:

    - name (optional): output variable name (all outputs when missing)
    - new_name (optional): optional renaming when `name` is defined
    - id (optional): node id
    - label (optional): used when `id` is missing
    - task_identifier (optional): used when `id` is missing
    - all (optional): used when `id`, `label` and `task_identifier` are missing (`True`: all nodes, `False`: end nodes)
    """
    if outputs is None:
        outputs = [{"all": False}]
    parsed = list()
    returned = {"id", "name", "new_name"}
    for output_item in outputs:
        if "id" in output_item:
            parsed.append({k: v for k, v in output_item.items() if k in returned})
            continue

        node_filters = dict()
        for k in ("label", "task_identifier"):
            if k in output_item:
                node_filters[k] = output_item[k]

        if node_filters:
            node_ids = iter_node_ids(graph, **node_filters)
        elif output_item.get("all"):
            node_ids = graph.nodes
        else:
            node_ids = end_nodes(graph)

        for node_id in node_ids:
            output_item = {k: v for k, v in output_item.items() if k in returned}
            output_item["id"] = node_id
            parsed.append(output_item)

    return parsed





def iter_node_ids(
    graph: networkx.DiGraph,
    label: Optional[str] = None,
    task_identifier: Optional[str] = None,
) -> Iterator[NodeIdType]:
    """Yield nodes with matching `label` AND `task_identifier`"""
    for node_id, node_attrs in graph.nodes.items():
        return_id = False
        if label is not None:
            node_label = get_node_label(node_id, node_attrs)
            if label != node_label:
                continue
            return_id = True
        if task_identifier is not None:
            s = node_attrs.get("task_identifier")
            if not s or not s.endswith(task_identifier):
                continue
            return_id = True
        if return_id:
            yield node_id





def extract_output_values(
    node_id: NodeIdType, task_or_outputs: Union[Task, Mapping], outputs: List[dict]
) -> Optional[dict]:
    """Output items have the following keys:

    - id: node id
    - label (optional): used when `id` is missing
    - name (optional): output variable name (all outputs when missing)
    - new_name (optional): optional renaming when name is defined
    """
    output_values = None
    if isinstance(task_or_outputs, Task):
        task_output_values = None
    else:
        task_output_values = task_or_outputs
    for output_item in outputs:
        if output_item.get("id") != node_id:
            continue
        if task_output_values is None:
            task_output_values = task_or_outputs.get_output_values()
        if output_values is None:
            output_values = dict()
        name = output_item.get("name")
        if name:
            new_name = output_item.get("new_name", name)
            output_values[new_name] = task_output_values.get(
                name, missing_data.MISSING_DATA
            )
        else:
            output_values.update(task_output_values)
    return output_values





def add_output_values(
    output_values: dict,
    node_id: NodeIdType,
    task_or_outputs: Union[Task, Dict],
    outputs: List[dict],
    merge_outputs: Optional[bool] = True,
) -> None:
    """Output items have the following keys:

    - id: node id
    - label (optional): used when `id` is missing
    - name (optional): output variable name (all outputs when missing)
    - new_name (optional): optional renaming when name is defined
    """
    task_output_values = extract_output_values(node_id, task_or_outputs, outputs)
    if task_output_values is not None:
        if merge_outputs:
            output_values.update(task_output_values)
        else:
            output_values[node_id] = task_output_values