utils.py

from contextlib import contextmanager
from unittest.mock import patch

import meshroom
from meshroom.core import desc, pluginManager, loadPluginFolder
from meshroom.core.plugins import NodePlugin, NodePluginStatus

import os

@contextmanager
def registeredNodeTypes(nodeTypes: list[type[desc.Node]]):
    nodePluginsList = {}
    for nodeType in nodeTypes:
        nodePlugin = NodePlugin(nodeType)
        pluginManager.registerNode(nodePlugin)
        nodePluginsList[nodeType] = nodePlugin

    yield

    for nodeType in nodeTypes:
        pluginManager.unregisterNode(nodePluginsList[nodeType])


@contextmanager
def overrideNodeTypeVersion(nodeType: type[desc.Node], version: str):
    """ Helper context manager to override the version of a given node type. """
    unpatchedFunc = meshroom.core.nodeVersion
    with patch.object(
        meshroom.core,
        "nodeVersion",
        side_effect=lambda type: version if type is nodeType else unpatchedFunc(type),
    ):
        yield


def registerNodeDesc(nodeDesc: type[desc.Node]):
    name = nodeDesc.__name__
    if not pluginManager.isRegistered(name):
        pluginManager._nodePlugins[name] = NodePlugin(nodeDesc)


def unregisterNodeDesc(nodeDesc: type[desc.Node]):
    name = nodeDesc.__name__
    if pluginManager.isRegistered(name):
        del pluginManager._nodePlugins[name]


@contextmanager
def registeredPlugins(folder: str):
    plugins = loadPluginFolder(folder)

    yield

    for plugin in plugins:
        pluginManager.removePlugin(plugin)

@contextmanager
def overrideOsEnvironmentVariables(envVariables: dict):
    with patch.dict(os.environ, envVariables, clear=False):
        yield