Add itertools.PickleCache. (#2756)

src/gluonts/itertools.py

@@ -13,8 +13,11 @@
 
 import itertools
 import math
+import pickle
 import random
+import tempfile
 from dataclasses import dataclass, field
+from pathlib import Path
 from typing import (
     Callable,
     Dict,
@@ -154,6 +157,41 @@ def __len__(self) -> int:
         return len(self.iterable)
 
 
+@dataclass
+class PickleCached:
+    """A caching wrapper for ``iterable`` using ``pickle`` to store cached
+    values on disk.
+
+    See :class:`Cached` for more information.
+    """
+
+    iterable: SizedIterable
+    cached: bool = False
+    _path: Path = field(
+        default_factory=lambda: Path(
+            tempfile.NamedTemporaryFile(delete=False).name
+        )
+    )
+
+    def __iter__(self):
+        if not self.cached:
+            with open(self._path, "wb") as tmpfile:
+                for batch in batcher(self.iterable, 16):
+                    pickle.dump(batch, tmpfile)
+                    yield from batch
+            self.cached = True
+        else:
+            with open(self._path, "rb") as tmpfile:
+                while True:
+                    try:
+                        yield from pickle.load(tmpfile)
+                    except EOFError:
+                        return
+
+    def __del__(self):
+        self._path.unlink()
+
+
 @dataclass
 class PseudoShuffled:
     """

test/test_itertools.py

@@ -26,6 +26,7 @@
 from gluonts.itertools import (
     batcher,
     Cached,
+    PickleCached,
     Cyclic,
     IterableSlice,
     PseudoShuffled,
@@ -68,6 +69,7 @@ def test_pseudo_shuffled(data: Iterable) -> None:
     "data, expected_elements_per_iteration",
     [
         (Cached(range(4)), (list(range(4)),) * 5),
+        (PickleCached(range(4)), (list(range(4)),) * 5),
         (batcher(range(10), 3), ([[0, 1, 2], [3, 4, 5], [6, 7, 8], [9]], [])),
         (IterableSlice(range(10), 3), ([0, 1, 2],) * 5),
         (