1""" @package antlr3.dottreegenerator
2@brief ANTLR3 runtime package, tree module
3
4This module contains all support classes for AST construction and tree parsers.
5
6"""
7
8# begin[licence]
9#
10# [The "BSD licence"]
11# Copyright (c) 2005-2008 Terence Parr
12# All rights reserved.
13#
14# Redistribution and use in source and binary forms, with or without
15# modification, are permitted provided that the following conditions
16# are met:
17# 1. Redistributions of source code must retain the above copyright
18#    notice, this list of conditions and the following disclaimer.
19# 2. Redistributions in binary form must reproduce the above copyright
20#    notice, this list of conditions and the following disclaimer in the
21#    documentation and/or other materials provided with the distribution.
22# 3. The name of the author may not be used to endorse or promote products
23#    derived from this software without specific prior written permission.
24#
25# THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR
26# IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
27# OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
28# IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT,
29# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
30# NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
31# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
32# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
33# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
34# THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
35#
36# end[licence]
37
38# lot's of docstrings are missing, don't complain for now...
39# pylint: disable-msg=C0111
40
41from antlr3.tree import CommonTreeAdaptor
42import stringtemplate3
43
44class DOTTreeGenerator(object):
45    """
46    A utility class to generate DOT diagrams (graphviz) from
47    arbitrary trees.  You can pass in your own templates and
48    can pass in any kind of tree or use Tree interface method.
49    """
50
51    _treeST = stringtemplate3.StringTemplate(
52        template=(
53        "digraph {\n" +
54        "  ordering=out;\n" +
55        "  ranksep=.4;\n" +
56        "  node [shape=plaintext, fixedsize=true, fontsize=11, fontname=\"Courier\",\n" +
57        "        width=.25, height=.25];\n" +
58        "  edge [arrowsize=.5]\n" +
59        "  $nodes$\n" +
60        "  $edges$\n" +
61        "}\n")
62        )
63
64    _nodeST = stringtemplate3.StringTemplate(
65        template="$name$ [label=\"$text$\"];\n"
66        )
67
68    _edgeST = stringtemplate3.StringTemplate(
69        template="$parent$ -> $child$ // \"$parentText$\" -> \"$childText$\"\n"
70        )
71
72    def __init__(self):
73        ## Track node to number mapping so we can get proper node name back
74        self.nodeToNumberMap = {}
75
76        ## Track node number so we can get unique node names
77        self.nodeNumber = 0
78
79
80    def toDOT(self, tree, adaptor=None, treeST=_treeST, edgeST=_edgeST):
81        if adaptor is None:
82            adaptor = CommonTreeAdaptor()
83
84        treeST = treeST.getInstanceOf()
85
86        self.nodeNumber = 0
87        self.toDOTDefineNodes(tree, adaptor, treeST)
88
89        self.nodeNumber = 0
90        self.toDOTDefineEdges(tree, adaptor, treeST, edgeST)
91        return treeST
92
93
94    def toDOTDefineNodes(self, tree, adaptor, treeST, knownNodes=None):
95        if knownNodes is None:
96            knownNodes = set()
97
98        if tree is None:
99            return
100
101        n = adaptor.getChildCount(tree)
102        if n == 0:
103            # must have already dumped as child from previous
104            # invocation; do nothing
105            return
106
107        # define parent node
108        number = self.getNodeNumber(tree)
109        if number not in knownNodes:
110            parentNodeST = self.getNodeST(adaptor, tree)
111            treeST.setAttribute("nodes", parentNodeST)
112            knownNodes.add(number)
113
114        # for each child, do a "<unique-name> [label=text]" node def
115        for i in range(n):
116            child = adaptor.getChild(tree, i)
117
118            number = self.getNodeNumber(child)
119            if number not in knownNodes:
120                nodeST = self.getNodeST(adaptor, child)
121                treeST.setAttribute("nodes", nodeST)
122                knownNodes.add(number)
123
124            self.toDOTDefineNodes(child, adaptor, treeST, knownNodes)
125
126
127    def toDOTDefineEdges(self, tree, adaptor, treeST, edgeST):
128        if tree is None:
129            return
130
131        n = adaptor.getChildCount(tree)
132        if n == 0:
133            # must have already dumped as child from previous
134            # invocation; do nothing
135            return
136
137        parentName = "n%d" % self.getNodeNumber(tree)
138
139        # for each child, do a parent -> child edge using unique node names
140        parentText = adaptor.getText(tree)
141        for i in range(n):
142            child = adaptor.getChild(tree, i)
143            childText = adaptor.getText(child)
144            childName = "n%d" % self.getNodeNumber(child)
145            edgeST = edgeST.getInstanceOf()
146            edgeST.setAttribute("parent", parentName)
147            edgeST.setAttribute("child", childName)
148            edgeST.setAttribute("parentText", parentText)
149            edgeST.setAttribute("childText", childText)
150            treeST.setAttribute("edges", edgeST)
151            self.toDOTDefineEdges(child, adaptor, treeST, edgeST)
152
153
154    def getNodeST(self, adaptor, t):
155        text = adaptor.getText(t)
156        nodeST = self._nodeST.getInstanceOf()
157        uniqueName = "n%d" % self.getNodeNumber(t)
158        nodeST.setAttribute("name", uniqueName)
159        if text is not None:
160            text = text.replace('"', r'\\"')
161        nodeST.setAttribute("text", text)
162        return nodeST
163
164
165    def getNodeNumber(self, t):
166        try:
167            return self.nodeToNumberMap[t]
168        except KeyError:
169            self.nodeToNumberMap[t] = self.nodeNumber
170            self.nodeNumber += 1
171            return self.nodeNumber - 1
172
173
174def toDOT(tree, adaptor=None, treeST=DOTTreeGenerator._treeST, edgeST=DOTTreeGenerator._edgeST):
175    """
176    Generate DOT (graphviz) for a whole tree not just a node.
177    For example, 3+4*5 should generate:
178
179    digraph {
180        node [shape=plaintext, fixedsize=true, fontsize=11, fontname="Courier",
181            width=.4, height=.2];
182        edge [arrowsize=.7]
183        "+"->3
184        "+"->"*"
185        "*"->4
186        "*"->5
187    }
188
189    Return the ST not a string in case people want to alter.
190
191    Takes a Tree interface object.
192
193    Example of invokation:
194
195        import antlr3
196        import antlr3.extras
197
198        input = antlr3.ANTLRInputStream(sys.stdin)
199        lex = TLexer(input)
200        tokens = antlr3.CommonTokenStream(lex)
201        parser = TParser(tokens)
202        tree = parser.e().tree
203        print tree.toStringTree()
204        st = antlr3.extras.toDOT(t)
205        print st
206
207    """
208
209    gen = DOTTreeGenerator()
210    return gen.toDOT(tree, adaptor, treeST, edgeST)
211