Internal change

PiperOrigin-RevId: 806185951

Author: The TensorFlow Datasets Authors

tensorflow_datasets/core/load.py

@@ -18,6 +18,7 @@
 from __future__ import annotations
 
 from collections.abc import Iterable, Iterator, Mapping, Sequence
+import contextlib  # pylint: disable=unused-import
 import dataclasses
 import difflib
 import posixpath
@@ -657,12 +658,14 @@ def load(
       object documents the entire dataset, regardless of the `split` requested.
       Split-specific information is available in `ds_info.splits`.
   """  # fmt: skip
-  dbuilder = _fetch_builder(
-      name=name,
-      data_dir=data_dir,
-      builder_kwargs=builder_kwargs,
-      try_gcs=try_gcs,
-  )
+  ctx = contextlib.nullcontext()
+  with ctx:
+    dbuilder = _fetch_builder(
+        name=name,
+        data_dir=data_dir,
+        builder_kwargs=builder_kwargs,
+        try_gcs=try_gcs,
+    )
   _download_and_prepare_builder(dbuilder, download, download_and_prepare_kwargs)
 
   if as_dataset_kwargs is None:
@@ -827,12 +830,10 @@ def data_source(
   """  # fmt:skip
   builder_kwargs = _set_file_format_for_data_source(data_dir, builder_kwargs)
   _validate_file_format_for_data_source(builder_kwargs)
-  dbuilder = _fetch_builder(
-      name,
-      data_dir,
-      builder_kwargs,
-      try_gcs,
-  )
+
+  ctx = contextlib.nullcontext()
+  with ctx:
+    dbuilder = _fetch_builder(name, data_dir, builder_kwargs, try_gcs)
   _download_and_prepare_builder(dbuilder, download, download_and_prepare_kwargs)
   return dbuilder.as_data_source(
       split=split, decoders=decoders, deserialize_method=deserialize_method

tensorflow_datasets/testing/dataset_builder_testing.py

@@ -425,7 +425,10 @@ def _test_checksums(self):
         f"Some urls checksums are missing at: {filepath}\n{err_msg}",
     )
 
-  def _download_and_prepare_as_dataset(self, builder):
+  def _download_and_prepare_as_dataset(
+      self,
+      builder,
+  ):
     # Provide the manual dir only if builder has MANUAL_DOWNLOAD_INSTRUCTIONS
     # set.
 
@@ -461,7 +464,9 @@ def _download_and_prepare_as_dataset(self, builder):
           beam_runner=beam_runner,
       )
       with self._test_key_not_local_path(builder):
-        builder.download_and_prepare(download_config=download_config)
+        builder.download_and_prepare(
+            download_config=download_config,
+        )
 
     with self._subTest("as_dataset"):
       self._assertAsDataset(builder)