astopt.py

import ast
import inspect
from astpretty import pprint

def simple_name_equal(elt, target):
    if not (isinstance(elt, ast.Name) and isinstance(target, ast.Name)):
        return False
    return elt.id == target.id

class Optimizer(ast.NodeTransformer):
    def __call__(self, func):
        source = inspect.getsource(func)
        tree = ast.parse(source)
        tree = self.visit(tree)
        ast.fix_missing_locations(tree)

        frame = inspect.currentframe()
        exec(compile(tree, "<ast>", "exec"), frame.f_globals, frame.f_locals)
        return frame.f_locals.get(func.__name__)

    def visit_FunctionDef(self, node):
        node.decorator_list = list(filter(lambda decorator: decorator.id != "optimize", node.decorator_list))
        return self.generic_visit(node)
    
    def visit_ListComp(self, node):
        if len(node.generators) == 1:
            comp = node.generators[0]
            if isinstance(comp, ast.comprehension):
                if simple_name_equal(node.elt, comp.target):
                    print('Translating simple list comp to iterator unpacking into a list')
                    return ast.List(
                        elts = [
                            ast.Starred(
                                comp.iter,
                                ctx = ast.Load()
                            )
                        ],
                        ctx = ast.Load()
                    )
        return node
import timeit

optimize = Optimizer()

def test():
    result = sum([a for a in range(500)])
    result += sum([b for b in tuple(range(300))])
    return result

optimized_test = optimize(test)

assert test() == optimized_test()
normal_timing = timeit.timeit(stmt="test()", setup="from __main__ import test", number=50000)
print(normal_timing)
optimized_timing = timeit.timeit(stmt="optimized_test()", setup="from __main__ import optimized_test", number=50000)
print(optimized_timing, optimized_test)
print("Saved", normal_timing - optimized_timing, "seconds")