Coverage for src/mlstatpy/graph/graphviz_helper.py: 85%

62 statements  

« prev     ^ index     » next       coverage.py v6.4.1, created at 2022-06-13 20:42 +0200

1""" 

2@file 

3@brief graphviz helper 

4""" 

5import os 

6import sys 

7from pyquickhelper.loghelper import run_cmd 

8from pyquickhelper.helpgen.conf_path_tools import find_graphviz_dot 

9 

10 

11def run_graphviz(filename, image, engine="dot"): 

12 """ 

13 Run :epkg:`GraphViz`. 

14 

15 @param filename filename which contains the graph definition 

16 @param image output image 

17 @param engine *dot* or *neato* 

18 @return output of graphviz 

19 """ 

20 ext = os.path.splitext(image)[-1] 

21 if ext != ".png": 

22 raise Exception("extension should be .png not " + str(ext)) 

23 if sys.platform.startswith("win"): 

24 bin_ = os.path.dirname(find_graphviz_dot()) 

25 # if bin not in os.environ["PATH"]: 

26 # os.environ["PATH"] = os.environ["PATH"] + ";" + bin 

27 cmd = '"{0}\\{3}" -Tpng "{1}" -o "{2}"'.format( 

28 bin_, filename, image, engine) 

29 else: 

30 cmd = '"{0}" -Tpng "{1}" -o "{2}"'.format(engine, filename, image) 

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

32 if len(err) > 0: 

33 raise Exception( 

34 "Unable to run Graphviz\nCMD:\n{0}\nOUT:\n{1}\nERR:\n{2}".format(cmd, out, err)) 

35 return out 

36 

37 

38def edges2gv(vertices, edges): 

39 """ 

40 Converts a graph into a :epkg:`GraphViz` file format. 

41 

42 @param edges see below 

43 @param vertices see below 

44 @return gv format 

45 

46 The function creates a file ``<image>.gv``. 

47 

48 .. runpython:: 

49 :showcode: 

50 

51 from mlstatpy.graph.graphviz_helper import edges2gv 

52 gv = edges2gv([(1, "eee", "red")], 

53 [(1, 2, "blue"), (3, 4), (1, 3)]) 

54 print(gv) 

55 

56 """ 

57 memovertex = {} 

58 for v in vertices: 

59 if isinstance(v, tuple): 

60 if len(v) == 1: 

61 memovertex[v[0]] = None 

62 else: 

63 memovertex[v[0]] = v[1:] 

64 else: 

65 memovertex[v] = None 

66 for edge in edges: 

67 i, j = edge[:2] 

68 if i not in memovertex: 

69 memovertex[i] = None 

70 if j not in memovertex: 

71 memovertex[j] = None 

72 

73 li = ["digraph{"] 

74 for k, v in memovertex.items(): 

75 if v is None: 

76 li.append("%s ;" % k) 

77 elif len(v) == 1: 

78 li.append("\"%s\" [label=\"%s\"];" % (k, v[0])) 

79 elif len(v) == 2: 

80 li.append("\"%s\" [label=\"%s\",fillcolor=%s,color=%s];" % ( 

81 k, v[0], v[1], v[1])) 

82 else: 

83 raise ValueError("unable to understand " + str(v)) 

84 

85 for edge in edges: 

86 i, j = edge[:2] 

87 if len(edge) == 2: 

88 li.append("\"%s\" -> \"%s\";" % (i, j)) 

89 elif len(edge) == 3: 

90 li.append("\"%s\" -> \"%s\" [label=\"%s\"];" % (i, j, edge[2])) 

91 elif len(edge) == 4: 

92 li.append( 

93 "\"%s\" -> \"%s\" [label=\"%s\",color=%s];" % (i, j, edge[2], edge[3])) 

94 else: 

95 raise ValueError("unable to understand " + str(edge)) 

96 li.append("}") 

97 

98 text = "\n".join(li) 

99 return text 

100 

101 

102def draw_graph_graphviz(vertices, edges, image=None, engine="dot"): 

103 """ 

104 Draws a graph using :epkg:`Graphviz`. 

105 

106 @param edges see below 

107 @param vertices see below 

108 @param image output image, None, just returns the output 

109 @param engine *dot* or *neato* 

110 @return :epkg:`Graphviz` output or 

111 the dot text if *image* is None 

112 

113 The function creates a file ``<image>.gv`` if *image* is not None. 

114 :: 

115 

116 edges = [ (1,2, label, color), (3,4), (1,3), ... ] , liste d'arcs 

117 vertices = [ (1, label, color), (2), ... ] , liste de noeuds 

118 image = nom d'image (format png) 

119 

120 """ 

121 text = edges2gv(vertices, edges) 

122 if image is None: 

123 return text 

124 filename = image + ".gv" 

125 with open(filename, "w", encoding="utf-8") as f: 

126 f.write(text) 

127 

128 out = run_graphviz(filename, image, engine=engine) 

129 if not os.path.exists(image): 

130 raise FileNotFoundError( 

131 "GraphViz failed with no reason. '{0}' not found.".format(image)) 

132 return out