diff options
Diffstat (limited to 'Lib/pickletools.py')
| -rw-r--r-- | Lib/pickletools.py | 69 |
1 files changed, 45 insertions, 24 deletions
diff --git a/Lib/pickletools.py b/Lib/pickletools.py index 71c2aa1..6b86723 100644 --- a/Lib/pickletools.py +++ b/Lib/pickletools.py @@ -2282,40 +2282,61 @@ def genops(pickle): def optimize(p): 'Optimize a pickle string by removing unused PUT opcodes' - not_a_put = object() - gets = { not_a_put } # set of args used by a GET opcode - opcodes = [] # (startpos, stoppos, putid) + put = 'PUT' + get = 'GET' + oldids = set() # set of all PUT ids + newids = {} # set of ids used by a GET opcode + opcodes = [] # (op, idx) or (pos, end_pos) proto = 0 + protoheader = b'' for opcode, arg, pos, end_pos in _genops(p, yield_end_pos=True): if 'PUT' in opcode.name: - opcodes.append((pos, end_pos, arg)) + oldids.add(arg) + opcodes.append((put, arg)) + elif opcode.name == 'MEMOIZE': + idx = len(oldids) + oldids.add(idx) + opcodes.append((put, idx)) elif 'FRAME' in opcode.name: pass - else: - if 'GET' in opcode.name: - gets.add(arg) - elif opcode.name == 'PROTO': - assert pos == 0, pos + elif 'GET' in opcode.name: + if opcode.proto > proto: + proto = opcode.proto + newids[arg] = None + opcodes.append((get, arg)) + elif opcode.name == 'PROTO': + if arg > proto: proto = arg - opcodes.append((pos, end_pos, not_a_put)) - prevpos, prevarg = pos, None + if pos == 0: + protoheader = p[pos: end_pos] + else: + opcodes.append((pos, end_pos)) + else: + opcodes.append((pos, end_pos)) + del oldids # Copy the opcodes except for PUTS without a corresponding GET out = io.BytesIO() - opcodes = iter(opcodes) - if proto >= 2: - # Write the PROTO header before any framing - start, stop, _ = next(opcodes) - out.write(p[start:stop]) - buf = pickle._Framer(out.write) + # Write the PROTO header before any framing + out.write(protoheader) + pickler = pickle._Pickler(out, proto) if proto >= 4: - buf.start_framing() - for start, stop, putid in opcodes: - if putid in gets: - buf.commit_frame() - buf.write(p[start:stop]) - if proto >= 4: - buf.end_framing() + pickler.framer.start_framing() + idx = 0 + for op, arg in opcodes: + if op is put: + if arg not in newids: + continue + data = pickler.put(idx) + newids[arg] = idx + idx += 1 + elif op is get: + data = pickler.get(newids[arg]) + else: + data = p[op:arg] + pickler.framer.commit_frame() + pickler.write(data) + pickler.framer.end_framing() return out.getvalue() ############################################################################## |