[DataPipe] Properly cleanup unclosed files within generator function (#6997)

Co-authored-by: Philip Meier <github.pmeier@posteo.de>

torchvision/prototype/datasets/_builtin/celeba.py

@@ -30,6 +30,7 @@ class CelebACSVParser(IterDataPipe[Tuple[str, Dict[str, str]]]):
 
     def __iter__(self) -> Iterator[Tuple[str, Dict[str, str]]]:
         for _, file in self.datapipe:
+            try:
                 lines = (line.decode() for line in file)
 
                 if self.fieldnames:
@@ -47,7 +48,7 @@ class CelebACSVParser(IterDataPipe[Tuple[str, Dict[str, str]]]):
 
                 for line in csv.DictReader(lines, fieldnames=fieldnames, dialect="celeba"):
                     yield line.pop("image_id"), line
-
+            finally:
                 file.close()
 
 

torchvision/prototype/datasets/_builtin/mnist.py

@@ -37,6 +37,7 @@ class MNISTFileReader(IterDataPipe[torch.Tensor]):
 
     def __iter__(self) -> Iterator[torch.Tensor]:
         for _, file in self.datapipe:
+            try:
                 read = functools.partial(fromfile, file, byte_order="big")
 
                 magic = int(read(dtype=torch.int32, count=1))
@@ -56,7 +57,7 @@ class MNISTFileReader(IterDataPipe[torch.Tensor]):
 
                 for _ in range(stop - start):
                     yield read(dtype=dtype, count=count).reshape(shape)
-
+            finally:
                 file.close()
 
 

torchvision/prototype/datasets/_builtin/pcam.py

@@ -28,11 +28,12 @@ class PCAMH5Reader(IterDataPipe[Tuple[str, io.IOBase]]):
         import h5py
 
         for _, handle in self.datapipe:
+            try:
                 with h5py.File(handle) as data:
                     if self.key is not None:
                         data = data[self.key]
                     yield from data
-
+            finally:
                 handle.close()