Support len for SPDL dataloader

Differential Revision: D78133462

Pull Request resolved: #839

Author: vbourgin

src/spdl/source/utils.py

@@ -13,7 +13,7 @@
 import random
 import sys
 import time
-from collections.abc import Iterable, Iterator, Sequence
+from collections.abc import Iterable, Iterator, Sequence, Sized
 from typing import TypeVar
 
 from ._type import IterableWithShuffle
@@ -220,6 +220,14 @@ def __iter__(self) -> Iterator[T]:
         if self._shuffle_last:
             self._shuffle()
 
+    def __len__(self) -> int:
+        if isinstance(self.src, Sized):
+            return len(self.src)
+        else:
+            raise TypeError(
+                f"Source iterator of type {type(self.src)} does not support length"
+            )
+
 
 def embed_shuffle(
     src: IterableWithShuffle[T], /, *, shuffle_last: bool = False, epoch: int = 0