diff --git a/aidge_core/export_utils/export_registry.py b/aidge_core/export_utils/export_registry.py new file mode 100644 index 0000000000000000000000000000000000000000..cf0ba3b7078fab278b7dd215d7f497495a25bd64 --- /dev/null +++ b/aidge_core/export_utils/export_registry.py @@ -0,0 +1,71 @@ +from typing import Dict, Tuple, Set +import aidge_core +from aidge_core.export_utils import ExportNode + +from enum import Enum + +# Language +LANGUAGE = Enum('LANGUAGE', ['Cpp/C']) +# Define new type registry_key +registry_key = Tuple[str, aidge_core.dtype, aidge_core.dformat] + +# TODO: very naive implementation ! +# error handling should be added ! +class ExportLib(): # Should be abstract ? + """Aidge export lib, define a registry + """ + # Lib name usefull ? + _name:str = None + # Registry of exportNode + _export_node_registry:Dict[registry_key, ExportNode] = {} + # The language type usefull ? + _language: LANGUAGE = None + def __init__(self) -> None: + raise RuntimeError("ExportLib should not be instanciated") + @classmethod + def exportable(cls, node:aidge_core.Node)->bool: + """ + :param node: aidge_node to try to export + :type node: aidge_core.Node + :return: True if this node is exportable with this lib, else false + :rtype: bool + """ + # TODO: should return usable error that can be catch to know if only some keys have not been respected ! + key: registry_key = (node.type(), node.dtype(), node.dformat()) + return key in cls._export_node_registry + @classmethod + def supported_operators(cls)->Set[str]: + """ + :return: list of supported operator by this ExportLib + :rtype: List[str] + """ + operators = set() + for key in cls._export_node_registry.keys(): + operators.add(key[0]) + return operators + @classmethod + def get_export_node(cls, node:aidge_core.Node)->ExportNode: + """ + :param node: Node to transform + :type node: aidge_core.Node + :return: Corresponding export node. + :rtype: ExportNode + """ + if not cls.exportable(node): + raise ValueError("Node is not exportable ...") + key: registry_key = (node.type(), node.dtype(), node.dformat()) + return key in cls._export_node_registry + @classmethod + def add_export_node(cls, key:registry_key, eNode:ExportNode)->None: + cls._export_node_registry[key] = eNode + +def operator_register(lib: ExportLib, key:registry_key, *args): + """Helper decorator to register an :py:class:`ExportNode` to an :py:class:`ExportLib` + """ + def decorator(operator): + def wrapper(*args, **kwargs): + return operator(*args, **kwargs) + lib.add_export_node(key, operator) + return wrapper + return decorator +