Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/graphviz/jupyter_integration.py: 59%

1"""Display rendered graph as SVG in Jupyter Notebooks and QtConsole.""" 

2 

3import typing 

4 

5from . import piping 

6 

7__all__ = ['JUPYTER_FORMATS', 

8 'SUPPORTED_JUPYTER_FORMATS', 'DEFAULT_JUPYTER_FORMAT', 

9 'get_jupyter_format_mimetype', 

10 'JupyterIntegration'] 

11 

12_IMAGE_JPEG = 'image/jpeg' 

13 

14JUPYTER_FORMATS = {'jpeg': _IMAGE_JPEG, 

15 'jpg': _IMAGE_JPEG, 

16 'png': 'image/png', 

17 'svg': 'image/svg+xml'} 

18 

19SUPPORTED_JUPYTER_FORMATS = set(JUPYTER_FORMATS) 

20 

21DEFAULT_JUPYTER_FORMAT = next(_ for _ in SUPPORTED_JUPYTER_FORMATS if _ == 'svg') 

22 

23MIME_TYPES = {'image/jpeg': '_repr_image_jpeg', 

24 'image/png': '_repr_image_png', 

25 'image/svg+xml': '_repr_image_svg_xml'} 

26 

27assert MIME_TYPES.keys() == set(JUPYTER_FORMATS.values()) 

28 

29SVG_ENCODING = 'utf-8' 

30 

31 

32def get_jupyter_format_mimetype(jupyter_format: str) -> str: 

33 try: 

34 return JUPYTER_FORMATS[jupyter_format] 

35 except KeyError: 

36 raise ValueError(f'unknown jupyter_format: {jupyter_format!r}' 

37 f' (must be one of {sorted(JUPYTER_FORMATS)})') 

38 

39 

40def get_jupyter_mimetype_format(mimetype: str) -> str: 

41 if mimetype not in MIME_TYPES: 

42 raise ValueError(f'unsupported mimetype: {mimetype!r}' 

43 f' (must be one of {sorted(MIME_TYPES)})') 

44 

45 assert mimetype in JUPYTER_FORMATS.values() 

46 

47 for format, jupyter_mimetype in JUPYTER_FORMATS.items(): 

48 if jupyter_mimetype == mimetype: 

49 return format 

50 

51 raise RuntimeError # pragma: no cover 

52 

53 

54class JupyterIntegration(piping.Pipe): 

55 """Display rendered graph as SVG in Jupyter Notebooks and QtConsole.""" 

56 

57 _jupyter_mimetype = get_jupyter_format_mimetype(DEFAULT_JUPYTER_FORMAT) 

58 

59 def _repr_mimebundle_(self, 

60 include: typing.Optional[typing.Iterable[str]] = None, 

61 exclude: typing.Optional[typing.Iterable[str]] = None, 

62 **_) -> typing.Dict[str, typing.Union[bytes, str]]: 

63 r"""Return the rendered graph as IPython mimebundle. 

64 

65 Args: 

66 include: Iterable of mimetypes to include in the result. 

67 If not given or ``None``: ``['image/sxg+xml']``. 

68 exclude: Iterable of minetypes to exclude from the result. 

69 Overrides ``include``. 

70 

71 Returns: 

72 Mapping from mimetypes to data. 

73 

74 Example: 

75 >>> doctest_mark_exe() 

76 >>> import graphviz 

77 >>> dot = graphviz.Graph() 

78 >>> dot._repr_mimebundle_() # doctest: +ELLIPSIS 

79 {'image/svg+xml': '<?xml version=... 

80 >>> dot._repr_mimebundle_(include=['image/png']) # doctest: +ELLIPSIS 

81 {'image/png': b'\x89PNG... 

82 >>> dot._repr_mimebundle_(include=[]) 

83 {} 

84 >>> dot._repr_mimebundle_(include=['image/svg+xml', 'image/jpeg'], 

85 ... exclude=['image/svg+xml']) # doctest: +ELLIPSIS 

86 {'image/jpeg': b'\xff... 

87 >>> list(dot._repr_mimebundle_(include=['image/png', 'image/jpeg'])) 

88 ['image/jpeg', 'image/png'] 

89 

90 See also: 

91 IPython documentation: 

92 - https://ipython.readthedocs.io/en/stable/api/generated/IPython.display.html#functions 

93 - https://ipython.readthedocs.io/en/stable/config/integrating.html#MyObject._repr_mimebundle_ # noqa: E501 

94 - https://nbviewer.org/github/ipython/ipython/blob/master/examples/IPython%20Kernel/Custom%20Display%20Logic.ipynb#Custom-Mimetypes-with-_repr_mimebundle_ # noqa: E501 

95 """ 

96 include = set(include) if include is not None else {self._jupyter_mimetype} 

97 include -= set(exclude or []) 

98 return {mimetype: getattr(self, method_name)() 

99 for mimetype, method_name in MIME_TYPES.items() 

100 if mimetype in include} 

101 

102 def _repr_image_jpeg(self) -> bytes: 

103 """Return the rendered graph as JPEG bytes.""" 

104 return self.pipe(format='jpeg') 

105 

106 def _repr_image_png(self) -> bytes: 

107 """Return the rendered graph as PNG bytes.""" 

108 return self.pipe(format='png') 

109 

110 def _repr_image_svg_xml(self) -> str: 

111 """Return the rendered graph as SVG string.""" 

112 return self.pipe(format='svg', encoding=SVG_ENCODING)