def test_find_Compare(): expr = parse_expr('a < b < c == e + (f, (gh,))') found = find_names(expr) assert len(found) == 6 assert eq( found, [var('a'), var('b'), var('c'), var('e'), var('f'), var('gh')] )
def test_find_Compare(): expr = parse_expr('a < b < c == e + (f, (gh,))') found = find_names(expr) assert len(found) == 6 assert eq( found, [var('a'), var('b'), var('c'), var('e'), var('f'), var('gh')] )
def visit_ListComp(self, node): """Generate a curried lambda function [x + y for x, y in [[1, 4], [2, 5], [3, 6]]] becomes [[1, 4], [2, 5], [3, 6]]].map(([x, y]) => x + y) """ try: generator, = node.generators except ValueError: raise NotImplementedError( 'Only single loop comprehensions are allowed' ) names = find_names(generator.target) argslist = [ast.arg(arg=name.id, annotation=None) for name in names] if len(names) <= 1: signature = ast.arguments( args=argslist, vararg=None, kwonlyargs=[], kw_defaults=[], kwarg=None, defaults=[], ) else: signature = ast.List(elts=argslist, ctx=ast.Load()) array = generator.iter lam_sig = functools.partial(ast.Lambda, args=signature) filters = generator.ifs if filters: filt = ast.BoolOp(op=ast.And(), values=filters) # array.filter method = ast.Attribute(value=array, attr='filter', ctx=ast.Load()) # array.filter(func) array = ast.Call( func=method, args=[lam_sig(body=filt)], keywords=[] ) method = ast.Attribute(value=array, attr='map', ctx=ast.Load()) mapped = ast.Call( func=method, args=[lam_sig(body=node.elt)], keywords=[] ) result = self.visit(mapped) return result
def visit_ListComp(self, node): """Generate a curried lambda function [x + y for x, y in [[1, 4], [2, 5], [3, 6]]] becomes [[1, 4], [2, 5], [3, 6]]].map(([x, y]) => x + y) """ try: generator, = node.generators except ValueError: raise NotImplementedError( 'Only single loop comprehensions are allowed' ) names = find_names(generator.target) argslist = [ast.arg(arg=name.id, annotation=None) for name in names] if len(names) <= 1: signature = ast.arguments( args=argslist, vararg=None, kwonlyargs=[], kw_defaults=[], kwarg=None, defaults=[], ) else: signature = ast.List(elts=argslist, ctx=ast.Load()) array = generator.iter lam_sig = functools.partial(ast.Lambda, args=signature) filters = generator.ifs if filters: filt = ast.BoolOp(op=ast.And(), values=filters) # array.filter method = ast.Attribute(value=array, attr='filter', ctx=ast.Load()) # array.filter(func) array = ast.Call( func=method, args=[lam_sig(body=filt)], keywords=[] ) method = ast.Attribute(value=array, attr='map', ctx=ast.Load()) mapped = ast.Call( func=method, args=[lam_sig(body=node.elt)], keywords=[] ) result = self.visit(mapped) return result
def test_find_ListComp(): expr = parse_expr('[i for i in range(n) if i < 2]') found = find_names(expr) assert all(isinstance(f, ast.Name) for f in found) assert eq(found, [var('i'), store('i'), var('n')])
def test_find_Tuple(): expr = parse_expr('(a, (b, 1), (((c,),),))') found = find_names(expr) assert len(found) == 3 assert eq(found, [var('a'), var('b'), var('c')])
def test_find_Name(): expr = parse_expr('b') found = find_names(expr) assert len(found) == 1 assert eq(found[0], var('b'))
def test_find_dup_names(): expr = parse_expr('a + 1 * a') found = find_names(expr) assert len(found) == 1 assert eq(found[0], var('a'))
def test_find_ListComp(): expr = parse_expr('[i for i in range(n) if i < 2]') found = find_names(expr) assert all(isinstance(f, ast.Name) for f in found) assert eq(found, [var('i'), store('i'), var('n')])
def test_find_Tuple(): expr = parse_expr('(a, (b, 1), (((c,),),))') found = find_names(expr) assert len(found) == 3 assert eq(found, [var('a'), var('b'), var('c')])
def test_find_Name(): expr = parse_expr('b') found = find_names(expr) assert len(found) == 1 assert eq(found[0], var('b'))
def test_find_dup_names(): expr = parse_expr('a + 1 * a') found = find_names(expr) assert len(found) == 1 assert eq(found[0], var('a'))