1'''
2altgraph.Dot - Interface to the dot language
3============================================
4
5The :py:mod:`~altgraph.Dot` module provides a simple interface to the
6file format used in the `graphviz <http://www.research.att.com/sw/tools/graphviz/>`_
7program. The module is intended to offload the most tedious part of the process
8(the **dot** file generation) while transparently exposing most of its features.
9
10To display the graphs or to generate image files the `graphviz <http://www.research.att.com/sw/tools/graphviz/>`_
11package needs to be installed on the system, moreover the :command:`dot` and :command:`dotty` programs must
12be accesible in the program path so that they can be ran from processes spawned
13within the module.
14
15Example usage
16-------------
17
18Here is a typical usage::
19
20    from altgraph import Graph, Dot
21
22    # create a graph
23    edges = [ (1,2), (1,3), (3,4), (3,5), (4,5), (5,4) ]
24    graph = Graph.Graph(edges)
25
26    # create a dot representation of the graph
27    dot = Dot.Dot(graph)
28
29    # display the graph
30    dot.display()
31
32    # save the dot representation into the mydot.dot file
33    dot.save_dot(file_name='mydot.dot')
34
35    # save dot file as gif image into the graph.gif file
36    dot.save_img(file_name='graph', file_type='gif')
37
38Directed graph and non-directed graph
39-------------------------------------
40
41Dot class can use for both directed graph and non-directed graph
42by passing ``graphtype`` parameter.
43
44Example::
45
46    # create directed graph(default)
47    dot = Dot.Dot(graph, graphtype="digraph")
48
49    # create non-directed graph
50    dot = Dot.Dot(graph, graphtype="graph")
51
52Customizing the output
53----------------------
54
55The graph drawing process may be customized by passing
56valid :command:`dot` parameters for the nodes and edges. For a list of all
57parameters see the `graphviz <http://www.research.att.com/sw/tools/graphviz/>`_
58documentation.
59
60Example::
61
62    # customizing the way the overall graph is drawn
63    dot.style(size='10,10', rankdir='RL', page='5, 5' , ranksep=0.75)
64
65    # customizing node drawing
66    dot.node_style(1, label='BASE_NODE',shape='box', color='blue' )
67    dot.node_style(2, style='filled', fillcolor='red')
68
69    # customizing edge drawing
70    dot.edge_style(1, 2, style='dotted')
71    dot.edge_style(3, 5, arrowhead='dot', label='binds', labelangle='90')
72    dot.edge_style(4, 5, arrowsize=2, style='bold')
73
74
75.. note::
76
77   dotty (invoked via :py:func:`~altgraph.Dot.display`) may not be able to
78   display all graphics styles. To verify the output save it to an image file
79   and look at it that way.
80
81Valid attributes
82----------------
83
84    - dot styles, passed via the :py:meth:`Dot.style` method::
85
86        rankdir = 'LR'   (draws the graph horizontally, left to right)
87        ranksep = number (rank separation in inches)
88
89    - node attributes, passed via the :py:meth:`Dot.node_style` method::
90
91        style = 'filled' | 'invisible' | 'diagonals' | 'rounded'
92        shape = 'box' | 'ellipse' | 'circle' | 'point' | 'triangle'
93
94    - edge attributes, passed via the :py:meth:`Dot.edge_style` method::
95
96        style     = 'dashed' | 'dotted' | 'solid' | 'invis' | 'bold'
97        arrowhead = 'box' | 'crow' | 'diamond' | 'dot' | 'inv' | 'none' | 'tee' | 'vee'
98        weight    = number (the larger the number the closer the nodes will be)
99
100    - valid `graphviz colors <http://www.research.att.com/~erg/graphviz/info/colors.html>`_
101
102    - for more details on how to control the graph drawing process see the
103      `graphviz reference <http://www.research.att.com/sw/tools/graphviz/refs.html>`_.
104'''
105import os
106import warnings
107
108from altgraph import GraphError
109
110
111class Dot(object):
112    '''
113    A  class providing a **graphviz** (dot language) representation
114    allowing a fine grained control over how the graph is being
115    displayed.
116
117    If the :command:`dot` and :command:`dotty` programs are not in the current system path
118    their location needs to be specified in the contructor.
119    '''
120
121    def __init__(self, graph=None, nodes=None, edgefn=None, nodevisitor=None, edgevisitor=None, name="G", dot='dot', dotty='dotty', neato='neato', graphtype="digraph"):
122        '''
123        Initialization.
124        '''
125        self.name, self.attr = name, {}
126
127        assert graphtype in ['graph', 'digraph']
128        self.type = graphtype
129
130        self.temp_dot = "tmp_dot.dot"
131        self.temp_neo = "tmp_neo.dot"
132
133        self.dot, self.dotty, self.neato = dot, dotty, neato
134
135        # self.nodes: node styles
136        # self.edges: edge styles
137        self.nodes, self.edges = {}, {}
138
139        if graph is not None and nodes is None:
140            nodes = graph
141        if graph is not None and edgefn is None:
142            def edgefn(node, graph=graph):
143                return graph.out_nbrs(node)
144        if nodes is None:
145            nodes = ()
146
147        seen = set()
148        for node in nodes:
149            if nodevisitor is None:
150                style = {}
151            else:
152                style = nodevisitor(node)
153            if style is not None:
154                self.nodes[node] = {}
155                self.node_style(node, **style)
156                seen.add(node)
157        if edgefn is not None:
158            for head in seen:
159                for tail in (n for n in edgefn(head) if n in seen):
160                    if edgevisitor is None:
161                        edgestyle = {}
162                    else:
163                        edgestyle = edgevisitor(head, tail)
164                    if edgestyle is not None:
165                        if head not in self.edges:
166                            self.edges[head] = {}
167                        self.edges[head][tail] = {}
168                        self.edge_style(head, tail, **edgestyle)
169
170    def style(self, **attr):
171        '''
172        Changes the overall style
173        '''
174        self.attr = attr
175
176    def display(self, mode='dot'):
177        '''
178        Displays the current graph via dotty
179        '''
180
181        if  mode == 'neato':
182            self.save_dot(self.temp_neo)
183            neato_cmd = "%s -o %s %s" % (self.neato, self.temp_dot, self.temp_neo)
184            os.system(neato_cmd)
185        else:
186            self.save_dot(self.temp_dot)
187
188        plot_cmd = "%s %s" % (self.dotty, self.temp_dot)
189        os.system(plot_cmd)
190
191    def node_style(self, node, **kwargs):
192        '''
193        Modifies a node style to the dot representation.
194        '''
195        if node not in self.edges:
196            self.edges[node] = {}
197        self.nodes[node] = kwargs
198
199    def all_node_style(self, **kwargs):
200        '''
201        Modifies all node styles
202        '''
203        for node in self.nodes:
204            self.node_style(node, **kwargs)
205
206    def edge_style(self, head, tail, **kwargs):
207        '''
208        Modifies an edge style to the dot representation.
209        '''
210        if tail not in self.nodes:
211            raise GraphError("invalid node %s" % (tail,))
212
213        try:
214            if tail not in self.edges[head]:
215                self.edges[head][tail]= {}
216            self.edges[head][tail] = kwargs
217        except KeyError:
218            raise GraphError("invalid edge  %s -> %s " % (head, tail) )
219
220    def iterdot(self):
221        # write graph title
222        if self.type == 'digraph':
223            yield 'digraph %s {\n' % (self.name,)
224        elif self.type == 'graph':
225            yield 'graph %s {\n' % (self.name,)
226
227        else:
228            raise GraphError("unsupported graphtype %s" % (self.type,))
229
230        # write overall graph attributes
231        for attr_name, attr_value in sorted(self.attr.items()):
232            yield '%s="%s";' % (attr_name, attr_value)
233        yield '\n'
234
235        # some reusable patterns
236        cpatt  = '%s="%s",'      # to separate attributes
237        epatt  = '];\n'          # to end attributes
238
239        # write node attributes
240        for node_name, node_attr in sorted(self.nodes.items()):
241            yield '\t"%s" [' % (node_name,)
242            for attr_name, attr_value in sorted(node_attr.items()):
243                yield cpatt % (attr_name, attr_value)
244            yield epatt
245
246        # write edge attributes
247        for head in sorted(self.edges):
248            for tail in sorted(self.edges[head]):
249                if self.type == 'digraph':
250                    yield '\t"%s" -> "%s" [' % (head, tail)
251                else:
252                    yield '\t"%s" -- "%s" [' % (head, tail)
253                for attr_name, attr_value in sorted(self.edges[head][tail].items()):
254                    yield cpatt % (attr_name, attr_value)
255                yield epatt
256
257        # finish file
258        yield '}\n'
259
260    def __iter__(self):
261        return self.iterdot()
262
263    def save_dot(self, file_name=None):
264        '''
265        Saves the current graph representation into a file
266        '''
267
268        if not file_name:
269            warnings.warn(DeprecationWarning, "always pass a file_name")
270            file_name = self.temp_dot
271
272        fp   = open(file_name, "w")
273        try:
274            for chunk in self.iterdot():
275                fp.write(chunk)
276        finally:
277            fp.close()
278
279    def save_img(self, file_name=None, file_type="gif", mode='dot'):
280        '''
281        Saves the dot file as an image file
282        '''
283
284        if not file_name:
285            warnings.warn(DeprecationWarning, "always pass a file_name")
286            file_name = "out"
287
288        if  mode == 'neato':
289            self.save_dot(self.temp_neo)
290            neato_cmd = "%s -o %s %s" % (self.neato, self.temp_dot, self.temp_neo)
291            os.system(neato_cmd)
292            plot_cmd = self.dot
293        else:
294            self.save_dot(self.temp_dot)
295            plot_cmd = self.dot
296
297        file_name  = "%s.%s" % (file_name, file_type)
298        create_cmd = "%s -T%s %s -o %s" % (plot_cmd, file_type, self.temp_dot, file_name)
299        os.system(create_cmd)
300