PyDataset now implements __iter__. (#21330)

It can now be iterated without failure when it has a finite number of batches.

Previously, it was iterable by accident due to a legacy feature of Python that predates the introduction of `__iter__`. Specifically, Python would see `__getitem__` implemented and would iterate by passing sequential indices. However, `__getitem__` was expected to throw an `IndexError` to indicate the end of the iterator as `__len__` is ignored by Python.

Instead, we return an iterator that knows about the length of the dataset.

Fixes #21151

Author: hertschuh

keras/src/trainers/data_adapters/py_dataset_adapter.py

@@ -152,8 +152,24 @@ def __getitem__(self, index):
         Returns:
             A batch
         """
+        del index
         raise NotImplementedError
 
+    def __iter__(self):
+        index_range = None
+        try:
+            num_batches = self.num_batches
+            if num_batches is not None:
+                index_range = range(num_batches)
+        except NotImplementedError:
+            pass
+
+        if index_range is None:
+            index_range = itertools.count()
+
+        for index in index_range:
+            yield self[index]
+
     @property
     def num_batches(self):
         """Number of batches in the PyDataset.

keras/src/trainers/data_adapters/py_dataset_adapter_test.py

@@ -422,3 +422,32 @@ def test_exception_reported(
             expected_exception_class, "Expected exception"
         ):
             next(it)
+
+    def test_iterate_finite(self):
+        py_dataset = ExamplePyDataset(
+            np.ones((6, 11), dtype="int32"),
+            np.zeros((6, 11), dtype="int32"),
+            batch_size=2,
+        )
+        batches = [batch for batch in py_dataset]
+        self.assertLen(batches, 3)
+
+    def test_iterate_infinite_with_none_num_batches(self):
+        py_dataset = ExamplePyDataset(
+            np.ones((6, 11), dtype="int32"),
+            np.zeros((6, 11), dtype="int32"),
+            batch_size=2,
+            infinite=True,
+        )
+        for index, _ in enumerate(py_dataset):
+            if index >= 10:
+                break
+
+    def test_iterate_infinite_with_no_len(self):
+        class NoLenDataset(py_dataset_adapter.PyDataset):
+            def __getitem__(self, idx):
+                yield np.ones((2, 11), dtype="int32")
+
+        for index, _ in enumerate(NoLenDataset()):
+            if index >= 10:
+                break