Complete an open todo on pickletools -- add a pickle optimizer.
diff --git a/Lib/pickletools.py b/Lib/pickletools.py
index 98f80f1..ae02a36 100644
--- a/Lib/pickletools.py
+++ b/Lib/pickletools.py
@@ -10,9 +10,7 @@
Print a symbolic disassembly of a pickle.
'''
-__all__ = ['dis',
- 'genops',
- ]
+__all__ = ['dis', 'genops', 'optimize']
# Other ideas:
#
@@ -1858,6 +1856,33 @@
break
##############################################################################
+# A pickle optimizer.
+
+def optimize(p):
+ 'Optimize a pickle string by removing unused PUT opcodes'
+ gets = set() # set of args used by a GET opcode
+ puts = [] # (arg, startpos, stoppos) for the PUT opcodes
+ prevpos = None # set to pos if previous opcode was a PUT
+ for opcode, arg, pos in genops(p):
+ if prevpos is not None:
+ puts.append((prevarg, prevpos, pos))
+ prevpos = None
+ if 'PUT' in opcode.name:
+ prevarg, prevpos = arg, pos
+ elif 'GET' in opcode.name:
+ gets.add(arg)
+
+ # Copy the pickle string except for PUTS without a corresponding GET
+ s = []
+ i = 0
+ for arg, start, stop in puts:
+ j = stop if (arg in gets) else start
+ s.append(p[i:j])
+ i = stop
+ s.append(p[i:])
+ return ''.join(s)
+
+##############################################################################
# A symbolic pickle disassembler.
def dis(pickle, out=None, memo=None, indentlevel=4):