26
26
from types import FunctionType
27
27
from copyreg import dispatch_table
28
28
from copyreg import _extension_registry , _inverted_registry , _extension_cache
29
- from itertools import islice
29
+ from itertools import batched
30
30
from functools import partial
31
31
import sys
32
32
from sys import maxsize
@@ -1035,12 +1035,11 @@ def _batch_appends(self, items, obj):
1035
1035
1036
1036
it = iter (items )
1037
1037
start = 0
1038
- while True :
1039
- tmp = list (islice (it , self ._BATCHSIZE ))
1040
- n = len (tmp )
1038
+ for batch in batched (it , self ._BATCHSIZE ):
1039
+ n = len (batch )
1041
1040
if n > 1 :
1042
1041
write (MARK )
1043
- for i , x in enumerate (tmp , start ):
1042
+ for i , x in enumerate (batch , start ):
1044
1043
try :
1045
1044
save (x )
1046
1045
except BaseException as exc :
@@ -1049,14 +1048,11 @@ def _batch_appends(self, items, obj):
1049
1048
write (APPENDS )
1050
1049
elif n :
1051
1050
try :
1052
- save (tmp [0 ])
1051
+ save (
8000
batch [0 ])
1053
1052
except BaseException as exc :
1054
1053
exc .add_note (f'when serializing { _T (obj )} item { start } ' )
1055
1054
raise
1056
1055
write (APPEND )
1057
- # else tmp is empty, and we're done
1058
- if n < self ._BATCHSIZE :
1059
- return
1060
1056
start += n
1061
1057
1062
1058
def save_dict (self , obj ):
@@ -1087,12 +1083,11 @@ def _batch_setitems(self, items, obj):
1087
1083
return
1088
1084
1089
1085
it = iter (items )
1090
- while True :
1091
- tmp = list (islice (it , self ._BATCHSIZE ))
1092
- n = len (tmp )
1086
+ for batch in batched (it , self ._BATCHSIZE ):
1087
+ n = len (batch )
1093
1088
if n > 1 :
1094
1089
write (MARK )
1095
- for k , v in tmp :
1090
+ for k , v in batch :
1096
1091
save (k )
1097
1092
try :
1098
1093
save (v )
@@ -1101,17 +1096,14 @@ def _batch_setitems(self, items, obj):
1101
1096
raise
1102
1097
write (SETITEMS )
1103
1098
elif n :
1104
- k , v = tmp [0 ]
1099
+ k , v = batch [0 ]
1105
1100
save (k )
1106
1101
try :
1107
1102
save (v )
1108
1103
except BaseException as exc :
1109
1104
exc .add_note (f'when serializing { _T (obj )} item { k !r} ' )
1110
1105
raise
1111
1106
write (SETITEM )
1112
- # else tmp is empty, and we're done
1113
- if n < self ._BATCHSIZE :
1114
- return
1115
1107
1116
1108
def save_set (self , obj ):
1117
1109
save = self .save
@@ -1125,8 +1117,7 @@ def save_set(self, obj):
1125
1117
self .memoize (obj )
1126
1118
1127
1119
it = iter (obj )
1128
- while True :
1129
- batch = list (islice (it , self ._BATCHSIZE ))
1120
+ for batch in batched (it , self ._BATCHSIZE ):
1130
1121
n = len (batch )
1131
1122
if n > 0 :
1132
1123
write (MARK )
@@ -1137,8 +1128,6 @@ def save_set(self, obj):
1137
1128
exc .add_note (f'when serializing { _T (obj )} element' )
1138
1129
raise
1139
1130
write (ADDITEMS )
1140
- if n < self ._BATCHSIZE :
1141
- return
1142
1131
dispatch [set ] = save_set
1143
1132
1144
1133
def save_frozenset (self , obj ):
0 commit comments