113 lines
4.3 KiB
Python
113 lines
4.3 KiB
Python
"""Display rendered graph as SVG in Jupyter Notebooks and QtConsole."""
|
|
|
|
import typing
|
|
|
|
from . import piping
|
|
|
|
__all__ = ['JUPYTER_FORMATS',
|
|
'SUPPORTED_JUPYTER_FORMATS', 'DEFAULT_JUPYTER_FORMAT',
|
|
'get_jupyter_format_mimetype',
|
|
'JupyterIntegration']
|
|
|
|
_IMAGE_JPEG = 'image/jpeg'
|
|
|
|
JUPYTER_FORMATS = {'jpeg': _IMAGE_JPEG,
|
|
'jpg': _IMAGE_JPEG,
|
|
'png': 'image/png',
|
|
'svg': 'image/svg+xml'}
|
|
|
|
SUPPORTED_JUPYTER_FORMATS = set(JUPYTER_FORMATS)
|
|
|
|
DEFAULT_JUPYTER_FORMAT = next(_ for _ in SUPPORTED_JUPYTER_FORMATS if _ == 'svg')
|
|
|
|
MIME_TYPES = {'image/jpeg': '_repr_image_jpeg',
|
|
'image/png': '_repr_image_png',
|
|
'image/svg+xml': '_repr_image_svg_xml'}
|
|
|
|
assert MIME_TYPES.keys() == set(JUPYTER_FORMATS.values())
|
|
|
|
SVG_ENCODING = 'utf-8'
|
|
|
|
|
|
def get_jupyter_format_mimetype(jupyter_format: str) -> str:
|
|
try:
|
|
return JUPYTER_FORMATS[jupyter_format]
|
|
except KeyError:
|
|
raise ValueError(f'unknown jupyter_format: {jupyter_format!r}'
|
|
f' (must be one of {sorted(JUPYTER_FORMATS)})')
|
|
|
|
|
|
def get_jupyter_mimetype_format(mimetype: str) -> str:
|
|
if mimetype not in MIME_TYPES:
|
|
raise ValueError(f'unsupported mimetype: {mimetype!r}'
|
|
f' (must be one of {sorted(MIME_TYPES)})')
|
|
|
|
assert mimetype in JUPYTER_FORMATS.values()
|
|
|
|
for format, jupyter_mimetype in JUPYTER_FORMATS.items():
|
|
if jupyter_mimetype == mimetype:
|
|
return format
|
|
|
|
raise RuntimeError # pragma: no cover
|
|
|
|
|
|
class JupyterIntegration(piping.Pipe):
|
|
"""Display rendered graph as SVG in Jupyter Notebooks and QtConsole."""
|
|
|
|
_jupyter_mimetype = get_jupyter_format_mimetype(DEFAULT_JUPYTER_FORMAT)
|
|
|
|
def _repr_mimebundle_(self,
|
|
include: typing.Optional[typing.Iterable[str]] = None,
|
|
exclude: typing.Optional[typing.Iterable[str]] = None,
|
|
**_) -> typing.Dict[str, typing.Union[bytes, str]]:
|
|
r"""Return the rendered graph as IPython mimebundle.
|
|
|
|
Args:
|
|
include: Iterable of mimetypes to include in the result.
|
|
If not given or ``None``: ``['image/sxg+xml']``.
|
|
exclude: Iterable of minetypes to exclude from the result.
|
|
Overrides ``include``.
|
|
|
|
Returns:
|
|
Mapping from mimetypes to data.
|
|
|
|
Example:
|
|
>>> doctest_mark_exe()
|
|
>>> import graphviz
|
|
>>> dot = graphviz.Graph()
|
|
>>> dot._repr_mimebundle_() # doctest: +ELLIPSIS
|
|
{'image/svg+xml': '<?xml version=...
|
|
>>> dot._repr_mimebundle_(include=['image/png']) # doctest: +ELLIPSIS
|
|
{'image/png': b'\x89PNG...
|
|
>>> dot._repr_mimebundle_(include=[])
|
|
{}
|
|
>>> dot._repr_mimebundle_(include=['image/svg+xml', 'image/jpeg'],
|
|
... exclude=['image/svg+xml']) # doctest: +ELLIPSIS
|
|
{'image/jpeg': b'\xff...
|
|
>>> list(dot._repr_mimebundle_(include=['image/png', 'image/jpeg']))
|
|
['image/jpeg', 'image/png']
|
|
|
|
See also:
|
|
IPython documentation:
|
|
- https://ipython.readthedocs.io/en/stable/api/generated/IPython.display.html#functions
|
|
- https://ipython.readthedocs.io/en/stable/config/integrating.html#MyObject._repr_mimebundle_ # noqa: E501
|
|
- https://nbviewer.org/github/ipython/ipython/blob/master/examples/IPython%20Kernel/Custom%20Display%20Logic.ipynb#Custom-Mimetypes-with-_repr_mimebundle_ # noqa: E501
|
|
"""
|
|
include = set(include) if include is not None else {self._jupyter_mimetype}
|
|
include -= set(exclude or [])
|
|
return {mimetype: getattr(self, method_name)()
|
|
for mimetype, method_name in MIME_TYPES.items()
|
|
if mimetype in include}
|
|
|
|
def _repr_image_jpeg(self) -> bytes:
|
|
"""Return the rendered graph as JPEG bytes."""
|
|
return self.pipe(format='jpeg')
|
|
|
|
def _repr_image_png(self) -> bytes:
|
|
"""Return the rendered graph as PNG bytes."""
|
|
return self.pipe(format='png')
|
|
|
|
def _repr_image_svg_xml(self) -> str:
|
|
"""Return the rendered graph as SVG string."""
|
|
return self.pipe(format='svg', encoding=SVG_ENCODING)
|