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

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

2 

3from collections.abc import Iterable, Mapping, Set 

4from typing import Final 

5 

6from . import piping 

7 

8__all__ = ['JUPYTER_FORMATS', 

9 'SUPPORTED_JUPYTER_FORMATS', 'DEFAULT_JUPYTER_FORMAT', 

10 'get_jupyter_format_mimetype', 

11 'JupyterIntegration'] 

12 

13_IMAGE_JPEG = 'image/jpeg' 

14 

15JUPYTER_FORMATS: Final[Mapping[str, str]] = {'jpeg': _IMAGE_JPEG, 

16 'jpg': _IMAGE_JPEG, 

17 'png': 'image/png', 

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

19 

20SUPPORTED_JUPYTER_FORMATS: Final[Set[str]] = set(JUPYTER_FORMATS) 

21 

22DEFAULT_JUPYTER_FORMAT: Final = next(_ for _ in SUPPORTED_JUPYTER_FORMATS 

23 if _ == 'svg') 

24 

25MIME_TYPES: Final[Mapping[str, str]] = {'image/jpeg': '_repr_image_jpeg', 

26 'image/png': '_repr_image_png', 

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

28 

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

30 

31SVG_ENCODING: Final = 'utf-8' 

32 

33 

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

35 try: 

36 return JUPYTER_FORMATS[jupyter_format] 

37 except KeyError: 

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

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

40 

41 

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

43 if mimetype not in MIME_TYPES: 

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

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

46 

47 assert mimetype in JUPYTER_FORMATS.values() 

48 

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

50 if jupyter_mimetype == mimetype: 

51 return format 

52 

53 raise RuntimeError # pragma: no cover 

54 

55 

56class JupyterIntegration(piping.Pipe): 

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

58 

59 _jupyter_mimetype = get_jupyter_format_mimetype(DEFAULT_JUPYTER_FORMAT) 

60 

61 def _repr_mimebundle_(self, 

62 include: Iterable[str] | None = None, 

63 exclude: Iterable[str] | None = None, 

64 **_) -> dict[str, bytes | str]: 

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

66 

67 Args: 

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

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

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

71 Overrides ``include``. 

72 

73 Returns: 

74 Mapping from mimetypes to data. 

75 

76 Example: 

77 >>> doctest_mark_exe() 

78 >>> import graphviz 

79 >>> dot = graphviz.Graph() 

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

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

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

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

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

85 {} 

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

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

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

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

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

91 

92 See also: 

93 IPython documentation: 

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

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

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

97 """ 

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

99 include -= set(exclude or []) 

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

101 for mimetype, method_name in MIME_TYPES.items() 

102 if mimetype in include} 

103 

104 def _repr_image_jpeg(self) -> bytes: 

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

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

107 

108 def _repr_image_png(self) -> bytes: 

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

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

111 

112 def _repr_image_svg_xml(self) -> str: 

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

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