Source

astoptimizer / astoptimizer / __init__.py

Full commit
import sys
if sys.version_info < (2, 6):
    raise Exception("astoptimizer requires Python 2.6 or later")
import ast

UNSET = object()

_builtin_compile = compile

def parse_ast(code, filename=None, mode=None, flags=0, dont_inherit=False, optimize=-1):
    if not filename:
        filename = '<string>'
    if not mode:
        mode = "exec"
    flags |= ast.PyCF_ONLY_AST
    if sys.version_info >= (3, 2):
        return _builtin_compile(code, filename, mode, flags, dont_inherit, optimize)
    else:
        return _builtin_compile(code, filename, mode, flags, dont_inherit)

def compile_ast(tree, filename=None, mode=None, flags=0, dont_inherit=False, optimize=-1):
    """
    Compile AST to bytecode.
    """
    if not filename:
        filename = '<string>'
    if not mode:
        mode = "exec"
    if sys.version_info >= (3, 2):
        return _builtin_compile(tree, filename, mode, flags, dont_inherit, optimize)
    else:
        return _builtin_compile(tree, filename, mode, flags, dont_inherit)

from astoptimizer.config import Config
from astoptimizer.optimizer import Optimizer

def optimize_ast(tree, config=None):
    """
    Optimize an AST with the specified config (or with default config).
    Return the optimized AST
    """
    optimizer = Optimizer(config)
    return optimizer.optimize(tree)

def optimize_code(code, filename=None, mode=None, flags=0, dont_inherit=False, optimize=-1, config=None):
    """
    Compile code to bytecode using the AST optimizer with specified config
    (or default config).
    """
    tree = parse_ast(code, filename, mode, flags, dont_inherit, optimize)
    tree = optimize_ast(tree, config)
    return compile_ast(tree, filename, mode, flags, dont_inherit, optimize)

def patch_compile(config=None):
    """
    Hook the builtin compile() function to use the AST optimizer.
    """
    if config is None:
        config = Config()

    def compile(source, filename, mode, flags=0, dont_inherit=False, optimize=-1):
        tree = parse_ast(source, filename, mode, flags, dont_inherit, optimize)
        if flags & ast.PyCF_ONLY_AST:
            return tree
        tree = optimize_ast(tree, config)
        code = compile_ast(tree, filename, mode, flags, dont_inherit, optimize)
        return code

    __builtins__['compile'] = compile