mirrors/cpython
1
0
Fork 0
mirror of https://github.com/python/cpython.git synced 2025-09-27 10:50:04 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 13

gh-126317: Simplify pickle code by using itertools.batched() (GH-126323)

Browse source
This commit is contained in:
Lee Dong Wook 2024-11-02 23:07:32 +09:00 • committed by GitHub
parent 10eeec2d4f
commit bd4be5e67d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 22 additions and 39 deletions
Download patch file Download diff file Expand all files Collapse all files

45
Lib/pickle.py
View file

@ -26,7 +26,7 @@ Misc variables:
from types import FunctionType from types import FunctionType
from copyreg import dispatch_table from copyreg import dispatch_table
from copyreg import _extension_registry, _inverted_registry, _extension_cache from copyreg import _extension_registry, _inverted_registry, _extension_cache
from itertools import islice from itertools import batched
from functools import partial from functools import partial
import sys import sys
from sys import maxsize from sys import maxsize
@ -1033,31 +1033,26 @@ class _Pickler:
write(APPEND) write(APPEND)
return return
it = iter(items)
start = 0 start = 0
while True: for batch in batched(items, self._BATCHSIZE):
tmp = list(islice(it, self._BATCHSIZE)) batch_len = len(batch)
n = len(tmp) if batch_len != 1:
if n > 1:
write(MARK) write(MARK)
for i, x in enumerate(tmp, start): for i, x in enumerate(batch, start):
try: try:
save(x) save(x)
except BaseException as exc: except BaseException as exc:
exc.add_note(f'when serializing {_T(obj)} item {i}') exc.add_note(f'when serializing {_T(obj)} item {i}')
raise raise
write(APPENDS) write(APPENDS)
elif n: else:
try: try:
save(tmp[0]) save(batch[0])
except BaseException as exc: except BaseException as exc:
exc.add_note(f'when serializing {_T(obj)} item {start}') exc.add_note(f'when serializing {_T(obj)} item {start}')
raise raise
write(APPEND) write(APPEND)
# else tmp is empty, and we're done start += batch_len
if n < self._BATCHSIZE:
return
start += n
def save_dict(self, obj): def save_dict(self, obj):
if self.bin: if self.bin:
@ -1086,13 +1081,10 @@ class _Pickler:
write(SETITEM) write(SETITEM)
return return
it = iter(items) for batch in batched(items, self._BATCHSIZE):
while True: if len(batch) != 1:
tmp = list(islice(it, self._BATCHSIZE))
n = len(tmp)
if n > 1:
write(MARK) write(MARK)
for k, v in tmp: for k, v in batch:
save(k) save(k)
try: try:
save(v) save(v)
@ -1100,8 +1092,8 @@ class _Pickler:
exc.add_note(f'when serializing {_T(obj)} item {k!r}') exc.add_note(f'when serializing {_T(obj)} item {k!r}')
raise raise
write(SETITEMS) write(SETITEMS)
elif n: else:
k, v = tmp[0] k, v = batch[0]
save(k) save(k)
try: try:
save(v) save(v)
@ -1109,9 +1101,6 @@ class _Pickler:
exc.add_note(f'when serializing {_T(obj)} item {k!r}') exc.add_note(f'when serializing {_T(obj)} item {k!r}')
raise raise
write(SETITEM) write(SETITEM)
# else tmp is empty, and we're done
if n < self._BATCHSIZE:
return
def save_set(self, obj): def save_set(self, obj):
save = self.save save = self.save
@ -1124,11 +1113,7 @@ class _Pickler:
write(EMPTY_SET) write(EMPTY_SET)
self.memoize(obj) self.memoize(obj)
it = iter(obj) for batch in batched(obj, self._BATCHSIZE):
while True:
batch = list(islice(it, self._BATCHSIZE))
n = len(batch)
if n > 0:
write(MARK) write(MARK)
try: try:
for item in batch: for item in batch:
@ -1137,8 +1122,6 @@ class _Pickler:
exc.add_note(f'when serializing {_T(obj)} element') exc.add_note(f'when serializing {_T(obj)} element')
raise raise
write(ADDITEMS) write(ADDITEMS)
if n < self._BATCHSIZE:
return
dispatch[set] = save_set dispatch[set] = save_set
def save_frozenset(self, obj): def save_frozenset(self, obj):