Generates random Python functions using Markov chains
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

prettify.py 9.6 KiB

10 years ago
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314
  1. import opcode
  2. import sys
  3. import func_smash
  4. OP_LOAD = {"LOAD_CONST": True, "LOAD_FAST": False, "LOAD_GLOBAL": False}
  5. OP_BINARY = {"BINARY_POWER": "**", "BINARY_MULTIPLY": "*",
  6. "BINARY_DIVIDE": "/", "BINARY_MODULO": "%", "BINARY_ADD": "+",
  7. "BINARY_SUBTRACT": "-", "BINARY_FLOOR_DIVIDE": "//",
  8. "BINARY_TRUE_DIVIDE": "/", "BINARY_LSHIFT": "<<",
  9. "BINARY_RSHIFT": ">>", "BINARY_AND": "&", "BINARY_XOR": "^",
  10. "BINARY_OR": "|"}
  11. OP_INPLACE = {"INPLACE_FLOOR_DIVIDE": "//", "INPLACE_TRUE_DIVIDE": "/",
  12. "INPLACE_ADD": "+", "INPLACE_SUBTRACT": "-",
  13. "INPLACE_MULTIPLY": "*", "INPLACE_DIVIDE": "/",
  14. "INPLACE_MODULO": "%", "INPLACE_POWER": "**",
  15. "INPLACE_LSHIFT": "<<", "INPLACE_RSHIFT": ">>",
  16. "INPLACE_AND": "&", "INPLACE_XOR": "^", "INPLACE_OR": "|"}
  17. OP_SUBSCR = ("BINARY_SUBSCR", "INPLACE_SUBSCR")
  18. OP_BUILD = {"BUILD_TUPLE": ("(", ")"), "BUILD_LIST": ("[", "]"),
  19. "BUILD_SET": ("{", "}")}
  20. TAB = 4
  21. def prettify_function(func, indent=0):
  22. args = _get_func_args(func)
  23. _print(indent, "def {0}({1}):".format(func.func_name, args))
  24. prettify_code(func.__code__, indent=indent+TAB)
  25. def prettify_code(codeobj, indent=0):
  26. codes = []
  27. codestring = codeobj.co_code
  28. length = len(codestring)
  29. i = 0
  30. while i < length:
  31. code = ord(codestring[i])
  32. i += 1
  33. if code >= opcode.HAVE_ARGUMENT:
  34. arg = func_smash._get_argument(codeobj, codestring, i, code)
  35. i += 2
  36. codes.append((code, arg))
  37. else:
  38. codes.append((code, None))
  39. block = _parse_codestring(codes)
  40. block.display(indent)
  41. def _get_func_args(func):
  42. codeobj = func.__code__
  43. count = codeobj.co_argcount
  44. if count == 0:
  45. return ""
  46. return ", ".join([arg for arg in codeobj.co_varnames[:count]])
  47. def _parse_codestring(codes):
  48. stack = Stack()
  49. print_buffer = []
  50. main_block = block = Block()
  51. i2block = {}
  52. drops = []
  53. elses = []
  54. i = 0
  55. for instruction in codes:
  56. code, arg = instruction
  57. opname = opcode.opname[code]
  58. while i in elses:
  59. elses.remove(i)
  60. block.toggle()
  61. while i in drops:
  62. drops.remove(i)
  63. block = block.parent()
  64. i2block[i] = block
  65. if code >= opcode.HAVE_ARGUMENT:
  66. i += 3
  67. else:
  68. i += 1
  69. if opname in OP_LOAD:
  70. stack.push(arg, is_literal=OP_LOAD[opname])
  71. elif opname in OP_BINARY:
  72. tos, tos1 = stack.pop(), stack.pop()
  73. stack.push(" ".join((tos1, OP_BINARY[opname], tos)))
  74. elif opname in OP_INPLACE: # Works, but doesn't use inplace op magic
  75. tos, tos1 = stack.pop(), stack.pop()
  76. stack.push(" ".join((tos1, OP_INPLACE[opname], tos)))
  77. elif opname in OP_SUBSCR:
  78. tos, tos1 = stack.pop(), stack.pop()
  79. stack.push("{0}[{1}]".format(tos1, tos))
  80. elif opname in OP_BUILD:
  81. args = []
  82. for i in xrange(arg):
  83. args.append(stack.pop())
  84. args.reverse()
  85. start, end = OP_BUILD[opname]
  86. stack.push(start + ", ".join(args) + end)
  87. elif opname == "BUILD_MAP":
  88. stack.push("{}")
  89. elif opname == "STORE_FAST":
  90. block.put(arg, "=", stack.pop())
  91. elif opname == "STORE_MAP":
  92. key, value = stack.pop(), stack.pop()
  93. pair = ": ".join((key, value))
  94. oldmap = stack.pop()
  95. if oldmap == "{}":
  96. newmap = "{" + pair + "}"
  97. else:
  98. newmap = oldmap[:-1] + ", " + pair + "}"
  99. stack.push(newmap)
  100. elif opname == "LOAD_ATTR":
  101. tos = stack.pop()
  102. new_tos = tos + "." + arg
  103. stack.push(new_tos)
  104. elif opname == "POP_TOP":
  105. block.put(stack.pop())
  106. elif opname == "CALL_FUNCTION":
  107. numargs, numkwargs = arg
  108. args = []
  109. for i in xrange(numkwargs):
  110. value = stack.pop()
  111. key = stack.pop(never_literal=True)
  112. args.append("=".join((key, value)))
  113. for i in xrange(numargs):
  114. args.append(stack.pop())
  115. args.reverse()
  116. funcname = stack.pop()
  117. stack.push("{0}({1})".format(funcname, ", ".join(args)))
  118. elif opname == "PRINT_ITEM":
  119. print_buffer.append(stack.pop())
  120. elif opname == "PRINT_NEWLINE":
  121. block.put("print", ", ".join(print_buffer))
  122. print_buffer = []
  123. elif opname == "RETURN_VALUE":
  124. block.put("return", stack.pop())
  125. elif opname == "COMPARE_OP":
  126. tos, tos1 = stack.pop(), stack.pop()
  127. compare = " ".join((tos1, arg, tos))
  128. stack.push(compare)
  129. elif opname == "POP_JUMP_IF_FALSE":
  130. block = block.child()
  131. block.split(stack.pop())
  132. elses.append(arg)
  133. elif opname == "POP_JUMP_IF_TRUE":
  134. block = block.child()
  135. block.split(stack.pop())
  136. elses.append(arg)
  137. elif opname == "JUMP_ABSOLUTE":
  138. if arg < i:
  139. block = block.parent()
  140. else:
  141. drops.append(arg)
  142. elif opname == "JUMP_FORWARD":
  143. drops.append(arg + i)
  144. elif opname == "SETUP_LOOP":
  145. block = block.child(loop=True)
  146. elif opname == "POP_BLOCK":
  147. block = block.parent()
  148. else:
  149. raise NotImplementedError(opname, arg, stack)
  150. return main_block
  151. def _print(indentation, *args, **kwargs):
  152. argstring = " ".join([str(arg) for arg in args if str(arg)])
  153. if kwargs.get("debug"):
  154. # Ignore debug messages without -d flag
  155. if len(sys.argv) > 1 and sys.argv[1] == "-d":
  156. print " " * 50 + "#", argstring
  157. return
  158. print " " * indentation + argstring
  159. class Stack(object):
  160. def __init__(self):
  161. self._items = []
  162. def __iter__(self):
  163. while self._items:
  164. yield self.pop()
  165. def __repr__(self):
  166. s = reversed([repr(itm) if lit else itm for itm, lit in self._items])
  167. return "Stack[" + ", ".join(s) + "]"
  168. def push(self, item, is_literal=False):
  169. self._items.append((item, is_literal))
  170. def pop(self, never_literal=False):
  171. item, is_literal = self._items.pop()
  172. return repr(item) if is_literal and not never_literal else item
  173. class Block(object):
  174. def __init__(self, parent=None, loop=False):
  175. self._parent = parent
  176. self._loop = loop
  177. self._true_block = self._focus = []
  178. self._false_block = []
  179. self._split = None
  180. def __iter__(self):
  181. for item in self._render_lines():
  182. yield item[1]
  183. def __repr__(self):
  184. return str([item[1] for item in self._render_lines()])
  185. def _has_lines(self):
  186. return self._split or self._true_block or self._false_block
  187. def _render_subblock(self, block, lines, indent):
  188. for item in block:
  189. if isinstance(item, Block):
  190. lines += item._render_lines(indent)
  191. else:
  192. lines.append((indent, item))
  193. def _render_lines(self, indent=0):
  194. lines = []
  195. if self._split:
  196. lines.append((indent, (self._split,)))
  197. indent += TAB
  198. self._render_subblock(self._true_block, lines, indent)
  199. if self._false_block:
  200. lines.append((indent - TAB, ("else:",)))
  201. self._render_subblock(self._false_block, lines, indent)
  202. return lines
  203. def child(self, loop=False):
  204. if self._loop and not loop and not self._has_lines():
  205. loop = True
  206. child = Block(parent=self, loop=loop)
  207. self._focus.append(child)
  208. return child
  209. def parent(self):
  210. if self._parent:
  211. return self._parent
  212. raise RuntimeError("Popping an orphaned block")
  213. def put(self, *code):
  214. self._focus.append(code)
  215. def split(self, test):
  216. if self._loop:
  217. self._split = "while " + test + ":"
  218. else:
  219. self._split = "if " + test + ":"
  220. def toggle(self):
  221. if self._focus is self._true_block:
  222. self._focus = self._false_block
  223. else:
  224. self._focus = self._true_block
  225. def display(self, indent=0):
  226. for indent, code in self._render_lines(indent):
  227. _print(indent, *code)
  228. if __name__ == "__main__":
  229. def f1(a):
  230. b = a + 10 / 7.0
  231. d = long_func(a, b+9, c=42, d="43")
  232. print b, d
  233. ex_tuple = ("Hello", "world!", abcdef)
  234. ex_list = [1, 2, 3] * 3 + [4, 5, 6]
  235. ex_set = {99, 98, 97, 96, 95}
  236. ex_dict = {d: e, f: False, "testing": g}
  237. return ex_dict
  238. def f2(a, b, c):
  239. if cmp1:
  240. line1
  241. if cmp2:
  242. line2
  243. elif cmp3:
  244. line3
  245. elif cmp4:
  246. line4
  247. else:
  248. line5
  249. line6
  250. else:
  251. line7
  252. line8
  253. if cmp4:
  254. line9
  255. if cmp5:
  256. if cmp6:
  257. if cmp7:
  258. if cmp8:
  259. if cmp9:
  260. line10
  261. return line11
  262. def f3(x, y):
  263. if cmp1:
  264. while cmp1:
  265. if cmp2:
  266. while cmp3:
  267. line1
  268. else:
  269. while cmp4:
  270. line2
  271. line3
  272. line4
  273. if cmp5:
  274. line5
  275. return line6
  276. prettify_function(f1)
  277. print
  278. prettify_function(f2)
  279. print
  280. prettify_function(f3)