Add assert_is_not_blocked to as_dataset and as_numpy_iterator.

PiperOrigin-RevId: 652937007

Author: The TensorFlow Datasets Authors

tensorflow_datasets/core/dataset_builder.py

@@ -806,6 +806,8 @@ def as_data_source(
     Raises:
       NotImplementedError if the data was not generated using ArrayRecords.
     """
+    self.assert_is_not_blocked()
+
     # By default, return all splits
     if split is None:
       split = {s: s for s in self.info.splits}
@@ -951,6 +953,8 @@ def as_dataset(
       If `batch_size` is -1, will return feature dictionaries containing
       the entire dataset in `tf.Tensor`s instead of a `tf.data.Dataset`.
     """
+    self.assert_is_not_blocked()
+
     # pylint: enable=line-too-long
     if not self.data_path.exists():
       raise AssertionError(

tensorflow_datasets/core/dataset_builder_test.py

@@ -437,6 +437,25 @@ def test_assert_is_not_blocked(self):
         assert builder_3.assert_is_not_blocked()
       assert not_blocked_builder.assert_is_not_blocked() is None
 
+  def test_blocked_as_dataset_and_as_data_source(self):
+    for config, version, expected_msg in [
+        ("plus1", "0.0.1", "Version 0.0.1 is blocked"),
+        ("plus2", "0.0.2", "plus2 is blocked for version 0.0.2"),
+    ]:
+      with testing.tmp_dir(self.get_temp_dir()) as tmp_dir:
+        tmp_dir = epath.Path(tmp_dir)
+        blocked_builder = DummyDatasetWithBlockedVersions(
+            config=config, version=version, data_dir=tmp_dir
+        )
+        with pytest.raises(
+            utils.DatasetVariantBlockedError, match=expected_msg
+        ):
+          blocked_builder.as_dataset()
+        with pytest.raises(
+            utils.DatasetVariantBlockedError, match=expected_msg
+        ):
+          blocked_builder.as_data_source()
+
   def test_versioned_configs(self):
     with testing.tmp_dir(self.get_temp_dir()) as tmp_dir:
       tmp_dir = epath.Path(tmp_dir)