Add checks for TfdsDataSource to ensure the task is not blocked.

PiperOrigin-RevId: 653203660

Author: The TensorFlow Datasets Authors

tensorflow_datasets/core/dataset_builder.py

@@ -557,6 +557,15 @@ def is_prepared(self) -> bool:
     """Returns whether this dataset is already downloaded and prepared."""
     return self.data_path.exists()
 
+  def is_blocked(self) -> utils.IsBlocked:
+    """Returns whether this builder (version, config) is blocked."""
+    config_name = self.builder_config.name if self.builder_config else None
+    if blocked_versions := self.blocked_versions:
+      return blocked_versions.is_blocked(
+          version=self.version, config=config_name
+      )
+    return utils.IsBlocked(False)
+
   def assert_is_not_blocked(self) -> None:
     """Checks that the dataset is not blocked."""
     config_name = self.builder_config.name if self.builder_config else None

tensorflow_datasets/core/dataset_builder_test.py

@@ -401,6 +401,22 @@ def test_builder_configs_configs_with_multiple_versions(self):
         set(DummyDatasetWithVersionedConfigs.builder_configs.keys()),
     )
 
+  def test_is_blocked(self):
+    with testing.tmp_dir(self.get_temp_dir()) as tmp_dir:
+      tmp_dir = epath.Path(tmp_dir)
+      builder_1 = DummyDatasetWithBlockedVersions(
+          config="plus1", version="0.0.1", data_dir=tmp_dir
+      )
+      builder_2 = DummyDatasetWithBlockedVersions(
+          config="plus2", version="0.0.2", data_dir=tmp_dir
+      )
+      not_blocked_builder = DummyDatasetWithConfigs(
+          config="plus1", version="0.0.1", data_dir=tmp_dir
+      )
+      assert builder_1.is_blocked()
+      assert builder_2.is_blocked()
+      assert not not_blocked_builder.is_blocked()
+
   def test_assert_is_not_blocked(self):
     with testing.tmp_dir(self.get_temp_dir()) as tmp_dir:
       tmp_dir = epath.Path(tmp_dir)
@@ -736,7 +752,7 @@ def test_default(self):
     )
 
   def test_explicitly_passed(self):
-    # When a dir is explictly passed, use it.
+    # When a dir is explicitly passed, use it.
     self.assertBuildDataDir(
         self.builder._build_data_dir(self.other_data_dir), self.other_data_dir
     )
@@ -801,7 +817,7 @@ def test_default_multi_dir_duplicate(self):
     with self.assertRaisesRegex(ValueError, "found in more than one directory"):
       self.builder._build_data_dir(None)
 
-  def test_expicit_multi_dir(self):
+  def test_explicit_multi_dir(self):
     # If two data dirs contains the same version
     # Data dir is explicitly passed
     file_utils.add_data_dir(self.other_data_dir)
@@ -979,7 +995,7 @@ def test_update_dataset_info_keeps_data_source(
       assert len(info_proto.data_source_accesses) == 1
       assert info_proto.data_source_accesses[0].file_system.path == "/x/y"
       builder.download_and_prepare()
-      # Manually check information was indeed written in datset_info.json and
+      # Manually check information was indeed written in dataset_info.json and
       # can be reloaded:
       builder = testing.DummyMnist(data_dir=tmp_dir)
       info_proto = builder.info.as_proto

tensorflow_datasets/core/utils/__init__.py

@@ -86,4 +86,5 @@
 from tensorflow_datasets.core.utils.version import BlockedVersions
 from tensorflow_datasets.core.utils.version import DatasetVariantBlockedError
 from tensorflow_datasets.core.utils.version import Experiment
+from tensorflow_datasets.core.utils.version import IsBlocked
 from tensorflow_datasets.core.utils.version import Version

tensorflow_datasets/core/utils/version.py

@@ -50,9 +50,12 @@ class IsBlocked:
   blocked.
   """
 
-  result: bool
+  result: bool = False
   blocked_msg: str | None = None
 
+  def __bool__(self) -> bool:
+    return self.result
+
 
 @dataclasses.dataclass(frozen=True)
 class BlockedVersions:

tensorflow_datasets/core/utils/version_test.py

@@ -150,6 +150,12 @@ def test_str_to_version():
   assert v1 == v0
 
 
+@pytest.mark.parametrize('is_blocked_result', [True, False])
+def test_is_blocked_bool(is_blocked_result):
+  is_blocked = version.IsBlocked(is_blocked_result)
+  assert bool(is_blocked) == is_blocked_result
+
+
 @pytest.mark.parametrize(
     'blocked_version, blocked_config, expected_res, expected_msg',
     [