Source code for kedro_onnx.typing

"""Type aliases for Kedro ONNX plugin."""
from typing_extensions import Literal
from typing import Protocol, TypeVar
from onnx.onnx_ml_pb2 import GraphProto


[docs]OnnxFrameworks = Literal[
    'tensorflow',
    'sklearn',
    'keras',
    'sparkml',
    'coreml',
    'xgboost',
    'lightgbm',
    'catboost',
    'h2o',
    'onnx',
]
"""Literal for supported ONNX frameworks."""

[docs]IT = TypeVar('IT')
[docs]OT = TypeVar('OT')


[docs]class ModelProto(Protocol):
    """Protocol for ONNX model."""

[docs]    doc_string: str
[docs]    domain: str
[docs]    ir_version: int
[docs]    model_version: int
[docs]    producer_name: str
[docs]    producer_version: str
[docs]    graph: GraphProto

[docs]    def ParseFromString(self, s: bytes) -> None:
        """Parse serialized model.

        Args:
            s (bytes): Serialized model.
        """
        ...  # pragma: no cover

[docs]    def SerializeToString(self) -> bytes:
        """Serialize model.

        Returns:
            bytes: Serialized model.
        """
        ...  # pragma: no cover