visualize.py (7566B)
1 # -*- coding: utf-8 -*- 2 """ 3 State Machine Visualizer 4 ~~~~~~~~~~~~~~~~~~~~~~~~ 5 6 This code provides a module that can use graphviz to visualise the state 7 machines included in h2. These visualisations can be used as part of the 8 documentation of h2, and as a reference material to understand how the 9 state machines function. 10 11 The code in this module is heavily inspired by code in Automat, which can be 12 found here: https://github.com/glyph/automat. For details on the licensing of 13 Automat, please see the NOTICES.visualizer file in this folder. 14 15 This module is very deliberately not shipped with the rest of h2. This is 16 because it is of minimal value to users who are installing h2: its use 17 is only really for the developers of h2. 18 """ 19 import argparse 20 import collections 21 import sys 22 23 import graphviz 24 import graphviz.files 25 26 import h2.connection 27 import h2.stream 28 29 30 StateMachine = collections.namedtuple( 31 'StateMachine', ['fqdn', 'machine', 'states', 'inputs', 'transitions'] 32 ) 33 34 35 # This is all the state machines we currently know about and will render. 36 # If any new state machines are added, they should be inserted here. 37 STATE_MACHINES = [ 38 StateMachine( 39 fqdn='h2.connection.H2ConnectionStateMachine', 40 machine=h2.connection.H2ConnectionStateMachine, 41 states=h2.connection.ConnectionState, 42 inputs=h2.connection.ConnectionInputs, 43 transitions=h2.connection.H2ConnectionStateMachine._transitions, 44 ), 45 StateMachine( 46 fqdn='h2.stream.H2StreamStateMachine', 47 machine=h2.stream.H2StreamStateMachine, 48 states=h2.stream.StreamState, 49 inputs=h2.stream.StreamInputs, 50 transitions=h2.stream._transitions, 51 ), 52 ] 53 54 55 def quote(s): 56 return '"{}"'.format(s.replace('"', r'\"')) 57 58 59 def html(s): 60 return '<{}>'.format(s) 61 62 63 def element(name, *children, **attrs): 64 """ 65 Construct a string from the HTML element description. 66 """ 67 formatted_attributes = ' '.join( 68 '{}={}'.format(key, quote(str(value))) 69 for key, value in sorted(attrs.items()) 70 ) 71 formatted_children = ''.join(children) 72 return u'<{name} {attrs}>{children}</{name}>'.format( 73 name=name, 74 attrs=formatted_attributes, 75 children=formatted_children 76 ) 77 78 79 def row_for_output(event, side_effect): 80 """ 81 Given an output tuple (an event and its side effect), generates a table row 82 from it. 83 """ 84 point_size = {'point-size': '9'} 85 event_cell = element( 86 "td", 87 element("font", enum_member_name(event), **point_size) 88 ) 89 side_effect_name = ( 90 function_name(side_effect) if side_effect is not None else "None" 91 ) 92 side_effect_cell = element( 93 "td", 94 element("font", side_effect_name, **point_size) 95 ) 96 return element("tr", event_cell, side_effect_cell) 97 98 99 def table_maker(initial_state, final_state, outputs, port): 100 """ 101 Construct an HTML table to label a state transition. 102 """ 103 header = "{} -> {}".format( 104 enum_member_name(initial_state), enum_member_name(final_state) 105 ) 106 header_row = element( 107 "tr", 108 element( 109 "td", 110 element( 111 "font", 112 header, 113 face="menlo-italic" 114 ), 115 port=port, 116 colspan="2", 117 ) 118 ) 119 rows = [header_row] 120 rows.extend(row_for_output(*output) for output in outputs) 121 return element("table", *rows) 122 123 124 def enum_member_name(state): 125 """ 126 All enum member names have the form <EnumClassName>.<EnumMemberName>. For 127 our rendering we only want the member name, so we take their representation 128 and split it. 129 """ 130 return str(state).split('.', 1)[1] 131 132 133 def function_name(func): 134 """ 135 Given a side-effect function, return its string name. 136 """ 137 return func.__name__ 138 139 140 def build_digraph(state_machine): 141 """ 142 Produce a L{graphviz.Digraph} object from a state machine. 143 """ 144 digraph = graphviz.Digraph(node_attr={'fontname': 'Menlo'}, 145 edge_attr={'fontname': 'Menlo'}, 146 graph_attr={'dpi': '200'}) 147 148 # First, add the states as nodes. 149 seen_first_state = False 150 for state in state_machine.states: 151 if not seen_first_state: 152 state_shape = "bold" 153 font_name = "Menlo-Bold" 154 else: 155 state_shape = "" 156 font_name = "Menlo" 157 digraph.node(enum_member_name(state), 158 fontame=font_name, 159 shape="ellipse", 160 style=state_shape, 161 color="blue") 162 seen_first_state = True 163 164 # We frequently have vary many inputs that all trigger the same state 165 # transition, and only differ in terms of their input and side-effect. It 166 # would be polite to say that graphviz does not handle this very well. So 167 # instead we *collapse* the state transitions all into the one edge, and 168 # then provide a label that displays a table of all the inputs and their 169 # associated side effects. 170 transitions = collections.defaultdict(list) 171 for transition in state_machine.transitions.items(): 172 initial_state, event = transition[0] 173 side_effect, final_state = transition[1] 174 transition_key = (initial_state, final_state) 175 transitions[transition_key].append((event, side_effect)) 176 177 for n, (transition_key, outputs) in enumerate(transitions.items()): 178 this_transition = "t{}".format(n) 179 initial_state, final_state = transition_key 180 181 port = "tableport" 182 table = table_maker( 183 initial_state=initial_state, 184 final_state=final_state, 185 outputs=outputs, 186 port=port 187 ) 188 189 digraph.node(this_transition, 190 label=html(table), margin="0.2", shape="none") 191 192 digraph.edge(enum_member_name(initial_state), 193 '{}:{}:w'.format(this_transition, port), 194 arrowhead="none") 195 digraph.edge('{}:{}:e'.format(this_transition, port), 196 enum_member_name(final_state)) 197 198 return digraph 199 200 201 def main(): 202 """ 203 Renders all the state machines in h2 into images. 204 """ 205 program_name = sys.argv[0] 206 argv = sys.argv[1:] 207 208 description = """ 209 Visualize h2 state machines as graphs. 210 """ 211 epilog = """ 212 You must have the graphviz tool suite installed. Please visit 213 http://www.graphviz.org for more information. 214 """ 215 216 argument_parser = argparse.ArgumentParser( 217 prog=program_name, 218 description=description, 219 epilog=epilog 220 ) 221 argument_parser.add_argument( 222 '--image-directory', 223 '-i', 224 help="Where to write out image files.", 225 default=".h2_visualize" 226 ) 227 argument_parser.add_argument( 228 '--view', 229 '-v', 230 help="View rendered graphs with default image viewer", 231 default=False, 232 action="store_true" 233 ) 234 args = argument_parser.parse_args(argv) 235 236 for state_machine in STATE_MACHINES: 237 print(state_machine.fqdn, '...discovered') 238 239 digraph = build_digraph(state_machine) 240 241 if args.image_directory: 242 digraph.format = "png" 243 digraph.render(filename="{}.dot".format(state_machine.fqdn), 244 directory=args.image_directory, 245 view=args.view, 246 cleanup=True) 247 print(state_machine.fqdn, "...wrote image into", args.image_directory) 248 249 250 if __name__ == '__main__': 251 main()