make prototype datasets traversable (#4950)

Co-authored-by: Francisco Massa <fvsmassa@gmail.com>

test/test_prototype_builtin_datasets.py

@@ -2,6 +2,7 @@ import io
 
 import builtin_dataset_mocks
 import pytest
+from torch.utils.data.graph import traverse
 from torchdata.datapipes.iter import IterDataPipe
 from torchvision.prototype import datasets, features
 from torchvision.prototype.datasets._api import DEFAULT_DECODER
@@ -83,6 +84,10 @@ class TestCommon:
         if not any(isinstance(value, features.Feature) for value in sample.values()):
             raise AssertionError("The sample contained no feature.")
 
+    @dataset_parametrization()
+    def test_traversable(self, dataset, mock_info):
+        traverse(dataset)
+
 
 class TestQMNIST:
     @pytest.mark.parametrize(

torchvision/prototype/datasets/utils/_internal.py

 import enum
+import functools
 import gzip
 import io
 import lzma
@@ -101,35 +102,37 @@ class Enumerator(IterDataPipe[Tuple[int, D]]):
         yield from enumerate(self.datapipe, self.start)
 
 
-def getitem(*items: Any) -> Callable[[Any], Any]:
-    def wrapper(obj: Any) -> Any:
+def _getitem_closure(obj: Any, *, items: Tuple[Any, ...]) -> Any:
     for item in items:
         obj = obj[item]
     return obj
 
-    return wrapper
 
+def getitem(*items: Any) -> Callable[[Any], Any]:
+    return functools.partial(_getitem_closure, items=items)
 
-def path_accessor(getter: Union[str, Callable[[pathlib.Path], D]]) -> Callable[[Tuple[str, Any]], D]:
-    if isinstance(getter, str):
-        name = getter
 
-        def getter(path: pathlib.Path) -> D:
+def _path_attribute_accessor(path: pathlib.Path, *, name: str) -> D:
     return cast(D, getattr(path, name))
 
-    def wrapper(data: Tuple[str, Any]) -> D:
-        return getter(pathlib.Path(data[0]))  # type: ignore[operator]
 
-    return wrapper
+def _path_accessor_closure(data: Tuple[str, Any], *, getter: Callable[[pathlib.Path], D]) -> D:
+    return getter(pathlib.Path(data[0]))
 
 
-def path_comparator(getter: Union[str, Callable[[pathlib.Path], D]], value: D) -> Callable[[Tuple[str, Any]], bool]:
-    accessor = path_accessor(getter)
+def path_accessor(getter: Union[str, Callable[[pathlib.Path], D]]) -> Callable[[Tuple[str, Any]], D]:
+    if isinstance(getter, str):
+        getter = functools.partial(_path_attribute_accessor, name=getter)
+
+    return functools.partial(_path_accessor_closure, getter=getter)
+
 
-    def wrapper(data: Tuple[str, Any]) -> bool:
+def _path_comparator_closure(data: Tuple[str, Any], *, accessor: Callable[[Tuple[str, Any]], D], value: D) -> bool:
     return accessor(data) == value
 
-    return wrapper
+
+def path_comparator(getter: Union[str, Callable[[pathlib.Path], D]], value: D) -> Callable[[Tuple[str, Any]], bool]:
+    return functools.partial(_path_comparator_closure, accessor=path_accessor(getter), value=value)
 
 
 class CompressionType(enum.Enum):