import dis
import new

_cache = {}
def compile_expression(s):
    co = compile(s, '<check>', 'eval')
    try:
        checkop, co, expr = _cache[co.co_code]
    except KeyError:
        bytecode = co.co_code
        have_argument = chr(dis.HAVE_ARGUMENT)
        return_value = chr(dis.opmap['RETURN_VALUE'])
        i = 0
        while bytecode[i] != return_value:
            lasti = i
            if bytecode[i] >= have_argument:
                i += 2
            i += 1
        try:
            checkop, newtail, expr = optimized_table[bytecode[lasti:i]]
        except KeyError:
            return lambda x: x, co, "x"
        bytecode = bytecode[:lasti] + newtail
        newco = new.code(co.co_argcount, co.co_nlocals, co.co_stacksize,
                      co.co_flags, bytecode, co.co_consts, co.co_names,
                      co.co_varnames, co.co_filename, co.co_name,
                      co.co_firstlineno, co.co_lnotab, co.co_freevars,
                      co.co_cellvars)
        _cache[co.co_code] = checkop, newco, expr
        co = newco
    return checkop, co, expr

def make_optimized_table(operator_table):
    result = {}
    for key, expr in operator_table.items():
        if 'y' not in expr:
            arity = 1
            checkop = eval("lambda x: %s" % expr)
        else:
            arity = 2
            checkop = eval("lambda (x,y): %s" % expr)
        if isinstance(key, str):
            key = chr(dis.opmap[key])
        else:
            opname, oparg = key
            key = chr(dis.opmap[opname]) + chr(oparg & 0xFF) + chr(oparg >> 8)
        if arity > 1:
            value = chr(dis.opmap['BUILD_TUPLE']) + chr(arity) + chr(0)
        else:
            value = ''
        value += chr(dis.opmap['RETURN_VALUE'])
        result[key] = checkop, value, expr
    return result

optimized_table = make_optimized_table({
    ('COMPARE_OP', 0): ('x < y'),
    ('COMPARE_OP', 1): ('x <= y'),
    ('COMPARE_OP', 2): ('x == y'),
    ('COMPARE_OP', 3): ('x != y'), 
    ('COMPARE_OP', 4): ('x > y'), 
    ('COMPARE_OP', 5): ('x >= y'), 
    ('COMPARE_OP', 6): ('x in y'), 
    ('COMPARE_OP', 7): ('x not in y'), 
    ('COMPARE_OP', 8): ('x is y'), 
    ('COMPARE_OP', 9): ('x is not y'), 
    'UNARY_NOT':       ('not x'),
    })

def check(s):
    import sys
    frame = sys._getframe(-1)
    checkop, co, expr = compile_expression(s)
    try:
        args = eval(co, frame.f_globals, frame.f_locals)
    except Exception, e:
        print "evaluation failed: %s" % s
        print "orig-dis:"
        dis.dis(compile(s, '?', 'eval'))
        print "mangled-dis:"
        dis.dis(co)
        print e
    else:
        expr = expr.replace('x', '%s').replace('y', '%s')
        expr = expr % args 

        if not checkop(args):
            print "failed:", expr # checkop, args


if __name__ == '__main__':
    # example:
    def f():
        return 5
    def g():
        return 3
    check("f() * g() == 5")
    check("not f()")
    check("f() and g() or 0")
    check("f() == g()")
    i = 4
    check("i == f()")
    #check("len(f()) == 0")
    check("5 or not 6")
    check("5 and 0")
    #for i in range(1000):
    #    check("f() == 5")
    #import profile
    #profile.run("run()")
    #run()