diff options
Diffstat (limited to 'astroid/astpeephole.py')
-rw-r--r-- | astroid/astpeephole.py | 86 |
1 files changed, 86 insertions, 0 deletions
diff --git a/astroid/astpeephole.py b/astroid/astpeephole.py new file mode 100644 index 0000000..af03462 --- /dev/null +++ b/astroid/astpeephole.py @@ -0,0 +1,86 @@ +# copyright 2003-2015 LOGILAB S.A. (Paris, FRANCE), all rights reserved. +# contact http://www.logilab.fr/ -- mailto:contact@logilab.fr +# +# This file is part of astroid. +# +# astroid is free software: you can redistribute it and/or modify it +# under the terms of the GNU Lesser General Public License as published by the +# Free Software Foundation, either version 2.1 of the License, or (at your +# option) any later version. +# +# astroid is distributed in the hope that it will be useful, but +# WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or +# FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public License +# for more details. +# +# You should have received a copy of the GNU Lesser General Public License along +# with astroid. If not, see <http://www.gnu.org/licenses/>. +"""Small AST optimizations.""" + +import _ast + +from astroid import nodes + + +__all__ = ('ASTPeepholeOptimizer', ) + + +try: + _TYPES = (_ast.Str, _ast.Bytes) +except AttributeError: + _TYPES = (_ast.Str, ) + + +class ASTPeepholeOptimizer(object): + """Class for applying small optimizations to generate new AST.""" + + def optimize_binop(self, node): + """Optimize BinOps with string Const nodes on the lhs. + + This fixes an infinite recursion crash, where multiple + strings are joined using the addition operator. With a + sufficient number of such strings, astroid will fail + with a maximum recursion limit exceeded. The + function will return a Const node with all the strings + already joined. + Return ``None`` if no AST node can be obtained + through optimization. + """ + ast_nodes = [] + current = node + while isinstance(current, _ast.BinOp): + # lhs must be a BinOp with the addition operand. + if not isinstance(current.left, _ast.BinOp): + return + if (not isinstance(current.left.op, _ast.Add) + or not isinstance(current.op, _ast.Add)): + return + + # rhs must a str / bytes. + if not isinstance(current.right, _TYPES): + return + + ast_nodes.append(current.right.s) + current = current.left + + if (isinstance(current, _ast.BinOp) + and isinstance(current.left, _TYPES) + and isinstance(current.right, _TYPES)): + # Stop early if we are at the last BinOp in + # the operation + ast_nodes.append(current.right.s) + ast_nodes.append(current.left.s) + break + + if not ast_nodes: + return + + # If we have inconsistent types, bail out. + known = type(ast_nodes[0]) + if any(type(element) is not known + for element in ast_nodes[1:]): + return + + value = known().join(reversed(ast_nodes)) + newnode = nodes.Const(value) + return newnode |