Coverage for pyquickhelper/helpgen/graphviz_helper.py: 89%

1""" 

2@file 

3@brief Helper about graphviz. 

4""" 

5import os 

6from ..loghelper import run_cmd 

7from .conf_path_tools import find_graphviz_dot 

8 

9 

10def plot_graphviz(dot, ax=None, temp_dot=None, temp_img=None, dpi=300, 

11 keep_axis=False): 

12 """ 

13 Plots a dot graph into a :epkg:`matplotlib` plot. 

14 

15 :param dot: dot language 

16 :param ax: existing ax 

17 :param temp_dot: temporary file, if None, 

18 a file is created and removed 

19 :param temp_img: temporary image, if None, 

20 a file is created and removed 

21 :param dpi: dpi 

22 :param keep_axis: keep axis in the graph 

23 :return: ax 

24 """ 

25 if temp_dot is None: 

26 temp_dot = "temp_%d.dot" % id(dot) 

27 clean_dot = True 

28 else: 

29 clean_dot = False 

30 if temp_img is None: 

31 temp_img = "temp_%d.png" % id(dot) 

32 clean_img = True 

33 else: 

34 clean_img = False 

35 with open(temp_dot, "w", encoding="utf-8") as f: 

36 f.write(dot) 

37 dot_path = find_graphviz_dot() 

38 cmd = '"%s" -Gdpi=%d -Tpng -o "%s" "%s"' % ( 

39 dot_path, dpi, temp_img, temp_dot) 

40 out, err = run_cmd(cmd, wait=True) 

41 if err is not None: 

42 err = err.strip("\r\n\t ") 

43 if len(err) > 0: 

44 if clean_dot: 

45 os.remove(temp_dot) 

46 if clean_img and os.path.exists(temp_img): 

47 os.remove(temp_img) 

48 raise RuntimeError( # pragma: no cover 

49 "Unable to run command line" 

50 "\n---CMD---\n{}\n---OUT---\n{}" 

51 "\n---ERR---\n{}".format( 

52 cmd, out, err)) 

53 if ax is None: 

54 import matplotlib.pyplot as plt # pragma: no cover 

55 ax = plt.gca() # pragma: no cover 

56 image = plt.imread(temp_img) # pragma: no cover 

57 elif isinstance(ax, str) and ax == 'new': 

58 import matplotlib.pyplot as plt # pragma: no cover 

59 _, ax = plt.subplots(1, 1) # pragma: no cover 

60 image = plt.imread(temp_img) # pragma: no cover 

61 else: 

62 import matplotlib.pyplot as plt 

63 image = plt.imread(temp_img) 

64 ax.imshow(image) 

65 if not keep_axis: 

66 ax.set_axis_off() 

67 if clean_dot: 

68 os.remove(temp_dot) 

69 if clean_img and os.path.exists(temp_img): 

70 os.remove(temp_img) 

71 return ax