diff --git a/prettify.py b/prettify.py index 5affe8d..52842a6 100644 --- a/prettify.py +++ b/prettify.py @@ -3,11 +3,12 @@ import sys import func_smash -OP_HASLOAD = ("LOAD_FAST", "LOAD_CONST", "LOAD_GLOBAL") +OP_HASLOAD = {"LOAD_CONST": True, "LOAD_FAST": False, "LOAD_GLOBAL": False} OP_HASBINARY = {"BINARY_ADD": "+", "BINARY_SUBTRACT": "-", "BINARY_MULTIPLY": "*", "BINARY_DIVIDE": "/", "BINARY_POWER": "**", "BINARY_MODULO": "%"} -OP_HASBUILD = ("BUILD_TUPLE", "BUILD_LIST", "BUILD_MAP", "BUILD_SET") +OP_HASBUILD = {"BUILD_TUPLE": ("(", ")"), "BUILD_LIST": ("[", "]"), + "BUILD_SET": ("{", "}")} def prettify_function(func, indent=0): args = _get_func_args(func) @@ -38,39 +39,50 @@ def _print_codestring(codes, indent): opname = opcode.opname[code] _print(indent, opname, arg, "; stack ==", stack, debug=True) if opname in OP_HASLOAD: - stack.append(arg) + _push(stack, arg, literal=OP_HASLOAD[opname]) elif opname in OP_HASBINARY: - tos, tos1 = str(stack.pop()), str(stack.pop()) - stack.append(" ".join((tos1, OP_HASBINARY[opname], tos))) + tos, tos1 = _pop(stack), _pop(stack) + _push(stack, " ".join((tos1, OP_HASBINARY[opname], tos))) elif opname in OP_HASBUILD: args = [] for i in xrange(arg): - args.append(str(stack.pop())) - if opname == "BUILD_TUPLE": - stack.append("(" + ", ".join(args) + ")") + args.append(_pop(stack)) + args.reverse() + start, end = OP_HASBUILD[opname] + _push(stack, start + ", ".join(args) + end) + elif opname == "BUILD_MAP": + _push(stack, "{}") elif opname == "STORE_FAST": - _print(indent, arg, "=", stack.pop()) + _print(indent, arg, "=", _pop(stack)) + elif opname == "STORE_MAP": + pair = ": ".join((_pop(stack), _pop(stack))) + oldmap = _pop(stack) + if oldmap == "{}": + newmap = "{" + pair + "}" + else: + newmap = oldmap[:-1] + ", " + pair + "}" + _push(stack, newmap) elif opname == "POP_TOP": - _print(indent, stack.pop()) + _print(indent, _pop(stack)) elif opname == "CALL_FUNCTION": numargs, numkwargs = arg args = [] for i in xrange(numkwargs): - value = str(stack.pop()) - key = str(stack.pop()) + value = _pop(stack) + key = _pop(stack, never_literal=True) args.append("=".join((key, value))) for i in xrange(numargs): - args.append(str(stack.pop())) + args.append(_pop(stack)) args.reverse() - funcname = stack.pop() - stack.append("{0}({1})".format(funcname, ", ".join(args))) + funcname = _pop(stack) + _push(stack, "{0}({1})".format(funcname, ", ".join(args))) elif opname == "PRINT_ITEM": - print_buffer.append(stack.pop()) + print_buffer.append(_pop(stack)) elif opname == "PRINT_NEWLINE": _print(indent, "print", ", ".join(print_buffer)) print_buffer = [] elif opname == "RETURN_VALUE": - _print(indent, "return", stack.pop()) + _print(indent, "return", _pop(stack)) else: raise NotImplementedError(opname, arg, stack) @@ -81,22 +93,32 @@ def _get_func_args(func): return "" return ", ".join([arg for arg in codeobj.co_varnames[:count]]) +def _push(stack, item, literal=False): + stack.append((item, literal)) + +def _pop(stack, never_literal=False): + item, literal = stack.pop() + return repr(item) if literal and not never_literal else item + def _print(indentation, *args, **kwargs): argstring = " ".join([str(arg) for arg in args if str(arg)]) if kwargs.get("debug"): # Ignore debug messages without -d flag if len(sys.argv) > 1 and sys.argv[1] == "-d": - print " " * 40 + "#", argstring + print " " * 50 + "#", argstring return print " " * indentation + argstring if __name__ == "__main__": def f1(a): b = a + 10 / 7.0 - d = long_func(a, b+9, c=42, d=43) + d = long_func(a, b+9, c=42, d="43") print b, d - v, z - return d % 10 + ex_tuple = ("Hello", "world!", abcdef) + ex_list = [1, 2, 3] * 3 + [4, 5, 6] + ex_set = {99, 98, 97, 96, 95} + ex_dict = {d: e, f: False, "testing": g} + return ex_dict def f2(): pass