/ftplugin/python/logilab/astng/utils.py

  1# This program is free software; you can redistribute it and/or modify
  2# it under the terms of the GNU Lesser General Public License as published by
  3# the Free Software Foundation; either version 2 of the License, or
  4# (at your option) any later version.
  5#
  6# This program is distributed in the hope that it will be useful, but WITHOUT
  7# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
  8# FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public License for more details.
  9#
 10# You should have received a copy of the GNU Lesser General Public License along with
 11# this program; if not, write to the Free Software Foundation, Inc.,
 12# 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.
 13# copyright 2003-2010 LOGILAB S.A. (Paris, FRANCE), all rights reserved.
 14# contact http://www.logilab.fr/ -- mailto:contact@logilab.fr
 15# copyright 2003-2010 Sylvain Thenault, all rights reserved.
 16# contact mailto:thenault@gmail.com
 17#
 18# This file is part of logilab-astng.
 19#
 20# logilab-astng is free software: you can redistribute it and/or modify it
 21# under the terms of the GNU Lesser General Public License as published by the
 22# Free Software Foundation, either version 2.1 of the License, or (at your
 23# option) any later version.
 24#
 25# logilab-astng is distributed in the hope that it will be useful, but
 26# WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 27# FITNESS FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public License
 28# for more details.
 29#
 30# You should have received a copy of the GNU Lesser General Public License along
 31# with logilab-astng. If not, see <http://www.gnu.org/licenses/>.
 32"""this module contains some utilities to navigate in the tree or to
 33extract information from it
 34
 35"""
 36
 37__docformat__ = "restructuredtext en"
 38
 39from logilab.astng.exceptions import ASTNGBuildingException
 40
 41
 42class ASTWalker:
 43    """a walker visiting a tree in preorder, calling on the handler:
 44
 45    * visit_<class name> on entering a node, where class name is the class of
 46    the node in lower case
 47
 48    * leave_<class name> on leaving a node, where class name is the class of
 49    the node in lower case
 50    """
 51
 52    def __init__(self, handler):
 53        self.handler = handler
 54        self._cache = {}
 55
 56    def walk(self, node, _done=None):
 57        """walk on the tree from <node>, getting callbacks from handler"""
 58        if _done is None:
 59            _done = set()
 60        if node in _done:
 61            raise AssertionError((id(node), node, node.parent))
 62        _done.add(node)
 63        self.visit(node)
 64        for child_node in node.get_children():
 65            self.handler.set_context(node, child_node)
 66            assert child_node is not node
 67            self.walk(child_node, _done)
 68        self.leave(node)
 69        assert node.parent is not node
 70
 71    def get_callbacks(self, node):
 72        """get callbacks from handler for the visited node"""
 73        klass = node.__class__
 74        methods = self._cache.get(klass)
 75        if methods is None:
 76            handler = self.handler
 77            kid = klass.__name__.lower()
 78            e_method = getattr(handler, 'visit_%s' % kid,
 79                               getattr(handler, 'visit_default', None))
 80            l_method = getattr(handler, 'leave_%s' % kid,
 81                               getattr(handler, 'leave_default', None))
 82            self._cache[klass] = (e_method, l_method)
 83        else:
 84            e_method, l_method = methods
 85        return e_method, l_method
 86
 87    def visit(self, node):
 88        """walk on the tree from <node>, getting callbacks from handler"""
 89        method = self.get_callbacks(node)[0]
 90        if method is not None:
 91            method(node)
 92
 93    def leave(self, node):
 94        """walk on the tree from <node>, getting callbacks from handler"""
 95        method = self.get_callbacks(node)[1]
 96        if method is not None:
 97            method(node)
 98
 99
100class LocalsVisitor(ASTWalker):
101    """visit a project by traversing the locals dictionary"""
102    def __init__(self):
103        ASTWalker.__init__(self, self)
104        self._visited = {}
105
106    def visit(self, node):
107        """launch the visit starting from the given node"""
108        if node in self._visited:
109            return
110        self._visited[node] = 1 # FIXME: use set ?
111        methods = self.get_callbacks(node)
112        if methods[0] is not None:
113            methods[0](node)
114        if 'locals' in node.__dict__: # skip Instance and other proxy
115            for name, local_node in node.items():
116                self.visit(local_node)
117        if methods[1] is not None:
118            return methods[1](node)
119
120
121def _check_children(node):
122    """a helper function to check children - parent relations"""
123    for child in node.get_children():
124        ok = False
125        if child is None:
126            print "Hm, child of %s is None" % node
127            continue
128        if not hasattr(child, 'parent'):
129            print " ERROR: %s has child %s %x with no parent" % (node, child, id(child))
130        elif not child.parent:
131            print " ERROR: %s has child %s %x with parent %r" % (node, child, id(child), child.parent)
132        elif child.parent is not node:
133            print " ERROR: %s %x has child %s %x with wrong parent %s" % (node,
134                                      id(node), child, id(child), child.parent)
135        else:
136            ok = True
137        if not ok:
138            print "lines;", node.lineno, child.lineno
139            print "of module", node.root(), node.root().name
140            raise ASTNGBuildingException
141        _check_children(child)
142
143
144from _ast import PyCF_ONLY_AST
145def parse(string):
146    return compile(string, "<string>", 'exec', PyCF_ONLY_AST)
147
148class TreeTester(object):
149    '''A helper class to see _ast tree and compare with astng tree
150
151    indent: string for tree indent representation
152    lineno: bool to tell if we should print the line numbers
153
154    >>> tester = TreeTester('print')
155    >>> print tester.native_tree_repr()
156
157    <Module>
158    .   body = [
159    .   <Print>
160    .   .   nl = True
161    .   ]
162    >>> print tester.astng_tree_repr()
163    Module()
164        body = [
165        Print()
166            dest = 
167            values = [
168            ]
169        ]
170    '''
171
172    indent = '.   '
173    lineno = False
174
175    def __init__(self, sourcecode):
176        self._string = ''
177        self.sourcecode = sourcecode
178        self._ast_node = None
179        self.build_ast()
180
181    def build_ast(self):
182        """build the _ast tree from the source code"""
183        self._ast_node = parse(self.sourcecode)
184
185    def native_tree_repr(self, node=None, indent=''):
186        """get a nice representation of the _ast tree"""
187        self._string = ''
188        if node is None:
189            node = self._ast_node
190        self._native_repr_tree(node, indent)
191        return self._string
192
193
194    def _native_repr_tree(self, node, indent, _done=None):
195        """recursive method for the native tree representation"""
196        from _ast import Load as _Load, Store as _Store, Del as _Del
197        from _ast import AST as Node
198        if _done is None:
199            _done = set()
200        if node in _done:
201            self._string += '\nloop in tree: %r (%s)' % (node,
202                                            getattr(node, 'lineno', None))
203            return
204        _done.add(node)
205        self._string += '\n' + indent +  '<%s>' % node.__class__.__name__
206        indent += self.indent
207        if not hasattr(node, '__dict__'):
208            self._string += '\n' + self.indent + " ** node has no __dict__ " + str(node)
209            return
210        node_dict = node.__dict__
211        if hasattr(node, '_attributes'):
212            for a in node._attributes:
213                attr = node_dict[a]
214                if attr is None:
215                    continue
216                if a in ("lineno", "col_offset") and not self.lineno:
217                    continue
218                self._string +='\n' +  indent + a + " = " + repr(attr)
219        for field in node._fields or ():
220            attr = node_dict[field]
221            if attr is None:
222                continue
223            if isinstance(attr, list):
224                if not attr:
225                    continue
226                self._string += '\n' + indent + field + ' = ['
227                for elt in attr:
228                    self._native_repr_tree(elt, indent, _done)
229                self._string += '\n' + indent + ']'
230                continue
231            if isinstance(attr, (_Load, _Store, _Del)):
232                continue
233            if isinstance(attr, Node):
234                self._string += '\n' + indent + field + " = "
235                self._native_repr_tree(attr, indent, _done)
236            else:
237                self._string += '\n' + indent + field + " = " + repr(attr)
238
239
240    def build_astng_tree(self):
241        """build astng tree from the _ast tree
242        """
243        from logilab.astng.builder import ASTNGBuilder
244        tree = ASTNGBuilder().string_build(self.sourcecode)
245        return tree
246
247    def astng_tree_repr(self, ids=False):
248        """build the astng tree and return a nice tree representation"""
249        mod = self.build_astng_tree()
250        return mod.repr_tree(ids)
251
252
253__all__ = ('LocalsVisitor', 'ASTWalker',)
254