Source code for mlprodict.onnxrt.doc.nb_helper

"""
Helpers for notebooks.


:githublink:`%|py|5`
"""
from IPython.core.magic import magics_class, line_magic
from jyquickhelper import RenderJsDot
from pyquickhelper.ipythonhelper import MagicCommandParser, MagicClassWithHelpers
from pyquickhelper.cli.cli_helper import create_cli_parser
from ..onnx_inference import OnnxInference


def onnxview(graph, recursive=False, local=False, add_rt_shapes=False):
    """
    Displays an :epkg:`ONNX` graph into a notebook.

    :param graph: filename, bytes, or :epkg:`onnx` graph.
    :param recursive: display subgraph
    :param local: use local path to javascript dependencies,
        recommanded option if used on :epkg:`MyBinder`)
    :param add_rt_shapes: add information about the shapes
        the runtime was able to find out,
        the runtime has to be `'python'`


    :githublink:`%|py|23`
    """
    sess = OnnxInference(graph, skip_run=not add_rt_shapes)
    dot = sess.to_dot(recursive=recursive, add_rt_shapes=add_rt_shapes)
    return RenderJsDot(dot, local=local)


@magics_class
class OnnxNotebook(MagicClassWithHelpers):

    """
    Defines magic commands to help with notebooks

    .. versionadded:: 1.1


    :githublink:`%|py|36`
    """

    @line_magic
    def onnxview(self, line):
        """
        Defines ``%onnxview``
        which displays an :epkg:`ONNX` graph.

        .. nbref::
            :title: onnxview

            The magic command ``%onnxview model_onnx`` is equivalent to function
            :func:`onnxview <mlprodict.onnxrt.doc.nb_helper.onnxview>`:

            ::

                onnx_view(model_onnx)

            It displays a visual representation of an :epkg:`ONNX` graph.



        :githublink:`%|py|56`
        """
        parser = self.get_parser(
            lambda: create_cli_parser(onnxview, cls=MagicCommandParser,
                                      positional={'graph'}), "onnxview")
        args = self.get_args(line, parser)

        if args is not None:
            res = onnxview(args.graph, recursive=args.recursive,
                           local=args.local, add_rt_shapes=args.add_rt_shapes)
            return res
        return None


def register_onnx_magics(ip=None):  # pragma: no cover
    """
    Register magics function, can be called from a notebook.

    :param      ip:      from ``get_ipython()``


    :githublink:`%|py|74`
    """
    if ip is None:
        from IPython import get_ipython
        ip = get_ipython()
    ip.register_magics(OnnxNotebook)