def build_dot()

in tensorwatch/model_graph/hiddenlayer/graph.py [0:0]

• 47 lines of code
• 14 McCabe index (conditional complexity)

    def build_dot(self, orientation):
        """Generate a GraphViz Dot graph.

        Returns a GraphViz Digraph object.
        """
        from graphviz import Digraph

        # Build GraphViz Digraph
        dot = Digraph()
        dot.attr("graph", 
                 bgcolor=self.theme["background_color"],
                 color=self.theme["outline_color"],
                 fontsize=self.theme["font_size"],
                 fontcolor=self.theme["font_color"],
                 fontname=self.theme["font_name"],
                 margin=self.theme["margin"],
                 rankdir=orientation,
                 pad=self.theme["padding"])
        dot.attr("node", shape="box", 
                 style="filled", margin="0,0",
                 fillcolor=self.theme["fill_color"],
                 color=self.theme["outline_color"],
                 fontsize=self.theme["font_size"],
                 fontcolor=self.theme["font_color"],
                 fontname=self.theme["font_name"])
        dot.attr("edge", style="solid", 
                 color=self.theme["outline_color"],
                 fontsize=self.theme["font_size"],
                 fontcolor=self.theme["font_color"],
                 fontname=self.theme["font_name"])

        for k, n in self.nodes.items():
            label = "<tr><td cellpadding='6'>{}</td></tr>".format(n.title)
            if n.caption:
                label += "<tr><td>{}</td></tr>".format(n.caption)
            if n.repeat > 1:
                label += "<tr><td align='right' cellpadding='2'>x{}</td></tr>".format(n.repeat)
            label = "<<table border='0' cellborder='0' cellpadding='0'>{}</table>>".\
                format(label)

            # figure out tooltip
            tooltips = set()
            if len(n.params):
                tooltips.update([str(n.params)])
            if n.combo_params:
                for params in n.combo_params:
                    if len(params):
                        tooltips.update([str(params)])
            
            # figure out shape and color
            layer_overrides = self.theme.get('layer_overrides', {})
            op_overrides = layer_overrides.get(n.op or n.name, {})

            dot.node(str(k), label, tooltip=html.escape(' '.join(tooltips)), **op_overrides)
        for a, b, label in self.edges:
            if isinstance(label, (list, tuple)):
                label = ' ' + "x".join([str(l or "?") for l in label])

            dot.edge(str(a), str(b), label)
        return dot