Home | History | Annotate | Download | only in antlr3 
      1 """ @package antlr3.dottreegenerator
      2 @brief ANTLR3 runtime package, tree module
      3 
      4 This 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 
     41 from antlr3.tree import CommonTreeAdaptor
     42 import stringtemplate3
     43 
     44 class 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 
    174 def 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