class CbrtPattern(Pattern):
# X ** .33333 => numpy.cbrt(X)
pattern = ast.BinOp(Placeholder(0), ast.Pow(), ast.Constant(1./3., None))
@staticmethod
def sub():
return ast.Call(
func=ast.Attribute(value=ast.Name(id=mangle('numpy'),
ctx=ast.Load(),
annotation=None,
type_comment=None),
attr="cbrt", ctx=ast.Load()),
args=[Placeholder(0)], keywords=[])
extra_imports = [ast.Import([ast.alias('numpy', mangle('numpy'))])]
lambda: ast.Call(func=ast.Attribute(value=ast.Name(
id='__builtin__', ctx=ast.Load(), annotation=None),
attr="xrange",
ctx=ast.Load()),
args=[
ast.BinOp(left=Placeholder(0),
op=ast.Sub(),
right=ast.Num(n=1)),
ast.Num(n=-1),
ast.Num(n=-1)
],
keywords=[])),
# X * X => X ** 2
(ast.BinOp(left=Placeholder(0), op=ast.Mult(), right=Placeholder(0)),
lambda: ast.BinOp(left=Placeholder(0), op=ast.Pow(), right=ast.Num(n=2))),
# a + "..." + b => "...".join((a, b))
(ast.BinOp(left=ast.BinOp(left=Placeholder(0),
op=ast.Add(),
right=ast.Str(Placeholder(1))),
op=ast.Add(),
right=Placeholder(2)),
lambda: ast.Call(func=ast.Attribute(
ast.Attribute(ast.Name('__builtin__', ast.Load(), None), 'str',
ast.Load()), 'join', ast.Load()),
args=[
ast.Str(Placeholder(1)),
ast.Tuple([Placeholder(0),
Placeholder(2)], ast.Load())
],
class Square(Transformation):
"""
Replaces **2 by a call to numpy.square.
>>> import gast as ast
>>> from pythran import passmanager, backend
>>> node = ast.parse('a**2')
>>> pm = passmanager.PassManager("test")
>>> _, node = pm.apply(Square, node)
>>> print pm.dump(backend.Python, node)
import numpy
numpy.square(a)
>>> node = ast.parse('numpy.power(a,2)')
>>> pm = passmanager.PassManager("test")
>>> _, node = pm.apply(Square, node)
>>> print pm.dump(backend.Python, node)
import numpy
numpy.square(a)
"""
POW_PATTERN = ast.BinOp(AST_any(), ast.Pow(), ast.Num(2))
POWER_PATTERN = ast.Call(
ast.Attribute(ast.Name('numpy', ast.Load(), None), 'power',
ast.Load()), [AST_any(), ast.Num(2)], [])
def __init__(self):
Transformation.__init__(self)
def replace(self, value):
self.update = self.need_import = True
return ast.Call(
ast.Attribute(ast.Name('numpy', ast.Load(), None), 'square',
ast.Load()), [value], [])
def visit_Module(self, node):
self.need_import = False
self.generic_visit(node)
if self.need_import:
importIt = ast.Import(names=[ast.alias(name='numpy', asname=None)])
node.body.insert(0, importIt)
return node
def expand_pow(self, node, n):
if n == 0:
return ast.Num(1)
elif n == 1:
return node
else:
node_square = self.replace(node)
node_pow = self.expand_pow(node_square, n >> 1)
if n & 1:
return ast.BinOp(node_pow, ast.Mult(), copy.deepcopy(node))
else:
return node_pow
def visit_BinOp(self, node):
self.generic_visit(node)
if ASTMatcher(Square.POW_PATTERN).search(node):
return self.replace(node.left)
elif isinstance(node.op, ast.Pow) and isinstance(node.right, ast.Num):
n = node.right.n
if int(n) == n and n > 0:
return self.expand_pow(node.left, n)
else:
return node
else:
return node
def visit_Call(self, node):
self.generic_visit(node)
if ASTMatcher(Square.POWER_PATTERN).search(node):
return self.replace(node.args[0])
else:
return node
def sub():
return ast.BinOp(left=Placeholder(0),
op=ast.Pow(),
right=ast.Constant(2, None))
def sub():
return ast.BinOp(Placeholder(0), ast.Pow(), Placeholder(1))
