Complete an open todo on pickletools -- add a pickle optimizer.

1 files changed, 28 insertions, 3 deletions

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 @@ dis(pickle, out=None, memo=None, indentlevel=4)
    Print a symbolic disassembly of a pickle.
 '''
 
-__all__ = ['dis',
-           'genops',
-          ]
+__all__ = ['dis', 'genops', 'optimize']
 
 # Other ideas:
 #
@@ -1858,6 +1856,33 @@ def genops(pickle):
             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):