from enum import Enum from typing import Any, TYPE_CHECKING, Union import meshroom from meshroom.core import Version from meshroom.core.attribute import Attribute, GroupAttribute, ListAttribute from meshroom.core.node import Node if TYPE_CHECKING: from meshroom.core.graph import Graph class GraphIO: """Centralize Graph file keys and IO version.""" __version__ = "2.0" class Keys(object): """File Keys.""" # Doesn't inherit enum to simplify usage (GraphIO.Keys.XX, without .value) Header = "header" NodesVersions = "nodesVersions" ReleaseVersion = "releaseVersion" FileVersion = "fileVersion" Graph = "graph" Template = "template" class Features(Enum): """File Features.""" Graph = "graph" Header = "header" NodesVersions = "nodesVersions" PrecomputedOutputs = "precomputedOutputs" NodesPositions = "nodesPositions" @staticmethod def getFeaturesForVersion(fileVersion: Union[str, Version]) -> tuple["GraphIO.Features", ...]: """Return the list of supported features based on a file version. Args: fileVersion (str, Version): the file version Returns: tuple of GraphIO.Features: the list of supported features """ if isinstance(fileVersion, str): fileVersion = Version(fileVersion) features = [GraphIO.Features.Graph] if fileVersion >= Version("1.0"): features += [ GraphIO.Features.Header, GraphIO.Features.NodesVersions, GraphIO.Features.PrecomputedOutputs, ] if fileVersion >= Version("1.1"): features += [GraphIO.Features.NodesPositions] return tuple(features) class GraphSerializer: """Standard Graph serializer.""" def __init__(self, graph: "Graph") -> None: self._graph = graph def serialize(self) -> dict: """ Serialize the Graph. """ return { GraphIO.Keys.Header: self.serializeHeader(), GraphIO.Keys.Graph: self.serializeContent(), } @property def nodes(self) -> list[Node]: return self._graph.nodes def serializeHeader(self) -> dict: """Build and return the graph serialization header. The header contains metadata about the graph, such as the: - version of the software used to create it. - version of the file format. - version of the nodes types used in the graph. - template flag. """ header: dict[str, Any] = {} header[GraphIO.Keys.ReleaseVersion] = meshroom.__version__ header[GraphIO.Keys.FileVersion] = GraphIO.__version__ header[GraphIO.Keys.NodesVersions] = self._getNodeTypesVersions() return header def _getNodeTypesVersions(self) -> dict[str, str]: """Get registered versions of each node types in `nodes`, excluding CompatibilityNode instances.""" nodeTypes = set([node.nodeDesc.__class__ for node in self.nodes if isinstance(node, Node)]) nodeTypesVersions = { nodeType.__name__: version for nodeType in nodeTypes if (version := meshroom.core.nodeVersion(nodeType)) is not None } # Sort them by name (to avoid random order changing from one save to another). return dict(sorted(nodeTypesVersions.items())) def serializeContent(self) -> dict: """Graph content serialization logic.""" return {node.name: self.serializeNode(node) for node in sorted(self.nodes, key=lambda n: n.name)} def serializeNode(self, node: Node) -> dict: """Node serialization logic.""" return node.toDict() class TemplateGraphSerializer(GraphSerializer): """Serializer for serializing a graph as a template.""" def serializeHeader(self) -> dict: header = super().serializeHeader() header[GraphIO.Keys.Template] = True return header def serializeNode(self, node: Node) -> dict: """Adapt node serialization to template graphs. Instead of getting all the inputs and internal attribute keys, only get the keys of the attributes whose value is not the default one. The output attributes, UIDs, parallelization parameters and internal folder are not relevant for templates, so they are explicitly removed from the returned dictionary. """ # For now, implemented as a post-process to update the default serialization. nodeData = super().serializeNode(node) inputKeys = list(nodeData["inputs"].keys()) internalInputKeys = [] internalInputs = nodeData.get("internalInputs", None) if internalInputs: internalInputKeys = list(internalInputs.keys()) for attrName in inputKeys: attribute = node.attribute(attrName) # check that attribute is not a link for choice attributes if attribute.isDefault and not attribute.isLink: del nodeData["inputs"][attrName] for attrName in internalInputKeys: attribute = node.internalAttribute(attrName) # check that internal attribute is not a link for choice attributes if attribute.isDefault and not attribute.isLink: del nodeData["internalInputs"][attrName] # If all the internal attributes are set to their default values, remove the entry if len(nodeData["internalInputs"]) == 0: del nodeData["internalInputs"] del nodeData["outputs"] del nodeData["uid"] del nodeData["internalFolder"] del nodeData["parallelization"] return nodeData class PartialGraphSerializer(GraphSerializer): """Serializer to serialize a partial graph (a subset of nodes).""" def __init__(self, graph: "Graph", nodes: list[Node]): super().__init__(graph) self._nodes = nodes @property def nodes(self) -> list[Node]: """Override to consider only the subset of nodes.""" return self._nodes def serializeNode(self, node: Node) -> dict: """Adapt node serialization to partial graph serialization.""" # NOTE: For now, implemented as a post-process to the default serialization. nodeData = super().serializeNode(node) # Override input attributes with custom serialization logic, to handle attributes # connected to nodes that are not in the list of nodes to serialize. for attributeName in nodeData["inputs"]: nodeData["inputs"][attributeName] = self._serializeAttribute(node.attribute(attributeName)) # Clear UID for non-compatibility nodes, as the custom attribute serialization # can be impacting the UID by removing connections to missing nodes. if not node.isCompatibilityNode: del nodeData["uid"] return nodeData def _serializeAttribute(self, attribute: Attribute) -> Any: """ Serialize `attribute` (recursively for list/groups) and deal with attributes being connected to nodes that are not part of the partial list of nodes to serialize. """ linkParam = attribute.getLinkParam() if linkParam is not None: # Use standard link serialization if upstream node is part of the serialization. if linkParam.node in self.nodes: return attribute.getExportValue() # Skip link serialization otherwise. # If part of a list, this entry can be discarded. if isinstance(attribute.root, ListAttribute): return None # Otherwise, return the default value for this attribute. return attribute.defaultValue() if isinstance(attribute, ListAttribute): # Recusively serialize each child of the ListAttribute, skipping those for which the attribute # serialization logic above returns None. return [ exportValue for child in attribute if (exportValue := self._serializeAttribute(child)) is not None ] if isinstance(attribute, GroupAttribute): # Recursively serialize each child of the group attribute. return {name: self._serializeAttribute(child) for name, child in attribute.value.items()} return attribute.getExportValue()