From 11eb1eeb2be3d1d064b5c9dbc23c88fdf3429b32 Mon Sep 17 00:00:00 2001
From: DongWook Lee <sh95119@gmail.com>
Date: Sat, 2 Nov 2024 16:28:26 +0900
Subject: [PATCH] gh-126317: Simplify stdlib code by using itertools.batched()

---
 Lib/pickle.py | 31 ++++++++++---------------------
 1 file changed, 10 insertions(+), 21 deletions(-)

diff --git a/Lib/pickle.py b/Lib/pickle.py
index ed8138beb908ee3..e0c27b90d4e5231 100644
--- a/Lib/pickle.py
+++ b/Lib/pickle.py
@@ -26,7 +26,7 @@
 from types import FunctionType
 from copyreg import dispatch_table
 from copyreg import _extension_registry, _inverted_registry, _extension_cache
-from itertools import islice
+from itertools import batched
 from functools import partial
 import sys
 from sys import maxsize
@@ -1035,12 +1035,11 @@ def _batch_appends(self, items, obj):
 
         it = iter(items)
         start = 0
-        while True:
-            tmp = list(islice(it, self._BATCHSIZE))
-            n = len(tmp)
+        for batch in batched(it, self._BATCHSIZE):
+            n = len(batch)
             if n > 1:
                 write(MARK)
-                for i, x in enumerate(tmp, start):
+                for i, x in enumerate(batch, start):
                     try:
                         save(x)
                     except BaseException as exc:
@@ -1049,14 +1048,11 @@ def _batch_appends(self, items, obj):
                 write(APPENDS)
             elif n:
                 try:
-                    save(tmp[0])
+                    save(batch[0])
                 except BaseException as exc:
                     exc.add_note(f'when serializing {_T(obj)} item {start}')
                     raise
                 write(APPEND)
-            # else tmp is empty, and we're done
-            if n < self._BATCHSIZE:
-                return
             start += n
 
     def save_dict(self, obj):
@@ -1087,12 +1083,11 @@ def _batch_setitems(self, items, obj):
             return
 
         it = iter(items)
-        while True:
-            tmp = list(islice(it, self._BATCHSIZE))
-            n = len(tmp)
+        for batch in batched(it, self._BATCHSIZE):
+            n = len(batch)
             if n > 1:
                 write(MARK)
-                for k, v in tmp:
+                for k, v in batch:
                     save(k)
                     try:
                         save(v)
@@ -1101,7 +1096,7 @@ def _batch_setitems(self, items, obj):
                         raise
                 write(SETITEMS)
             elif n:
-                k, v = tmp[0]
+                k, v = batch[0]
                 save(k)
                 try:
                     save(v)
@@ -1109,9 +1104,6 @@ def _batch_setitems(self, items, obj):
                     exc.add_note(f'when serializing {_T(obj)} item {k!r}')
                     raise
                 write(SETITEM)
-            # else tmp is empty, and we're done
-            if n < self._BATCHSIZE:
-                return
 
     def save_set(self, obj):
         save = self.save
@@ -1125,8 +1117,7 @@ def save_set(self, obj):
         self.memoize(obj)
 
         it = iter(obj)
-        while True:
-            batch = list(islice(it, self._BATCHSIZE))
+        for batch in batched(it, self._BATCHSIZE):
             n = len(batch)
             if n > 0:
                 write(MARK)
@@ -1137,8 +1128,6 @@ def save_set(self, obj):
                     exc.add_note(f'when serializing {_T(obj)} element')
                     raise
                 write(ADDITEMS)
-            if n < self._BATCHSIZE:
-                return
     dispatch[set] = save_set
 
     def save_frozenset(self, obj):