/
jupyter_integration.py
112 lines (84 loc) · 4.24 KB
/
jupyter_integration.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
"""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)