Add BlockedVersions to document which TFDS versions and configs should not be used.

PiperOrigin-RevId: 652441832

Author: The TensorFlow Datasets Authors

tensorflow_datasets/core/dataset_builder.py

@@ -229,6 +229,9 @@ class DatasetBuilder(registered.RegisteredDataset):
   # Example: `imported_builder_cls` function in `registered.py` module sets it.
   pkg_dir_path: Optional[epath.Path] = None
 
+  # Holds information on versions and configs that should not be used.
+  BLOCKED_VERSIONS: ClassVar[Optional[utils.BlockedVersions]] = None
+
   @classmethod
   def _get_pkg_dir_path(cls) -> epath.Path:
     """Returns class pkg_dir_path, infer it first if not set."""
@@ -396,6 +399,10 @@ def release_notes(self) -> Dict[str, str]:
     else:
       return self.RELEASE_NOTES
 
+  @property
+  def blocked_versions(self) -> utils.BlockedVersions | None:
+    return self.BLOCKED_VERSIONS
+
   @property
   def data_dir_root(self) -> epath.Path:
     """Returns the root directory where all TFDS datasets are stored.

tensorflow_datasets/core/utils/__init__.py

@@ -83,5 +83,6 @@
 from tensorflow_datasets.core.utils.tqdm_utils import tqdm
 from tensorflow_datasets.core.utils.tqdm_utils import TqdmStream
 from tensorflow_datasets.core.utils.type_utils import *
+from tensorflow_datasets.core.utils.version import BlockedVersions
 from tensorflow_datasets.core.utils.version import Experiment
 from tensorflow_datasets.core.utils.version import Version

tensorflow_datasets/core/utils/version.py

@@ -17,6 +17,7 @@
 
 from __future__ import annotations
 
+import dataclasses
 import enum
 import re
 from typing import List, Tuple, Union
@@ -31,6 +32,41 @@
 _VERSION_RESOLVED_REG = re.compile(_VERSION_TMPL.format(v=_NO_LEADING_ZEROS))
 
 
+# A dictionary of blocked versions or configs.
+# The key is a version or config string, the value is a short sentence
+# explaining why that version or config should not be used (or None).
+BlockedWithMsg = dict[str, str | None]
+
+
+@dataclasses.dataclass(frozen=True)
+class BlockedVersions:
+  """Holds information on versions and configs that should not be used.
+
+  Note that only complete versions can be blocked: wilcards in versions are not
+  supported.
+
+  versions: A dictionary of bad versions for which all configs should be
+    blocked.
+  configs: A mapping from versions to a dictionary of configs that should not be
+    used for that version.
+  """
+
+  versions: BlockedWithMsg = dataclasses.field(default_factory=dict)
+  configs: dict[str, BlockedWithMsg] = dataclasses.field(default_factory=dict)
+
+  def is_blocked(
+      self, version: str | Version, config: str | None = None
+  ) -> bool:
+    """Checks whether a version or config is blocked."""
+    if isinstance(version, Version):
+      version = str(version)
+    if version in self.versions:
+      return True
+    if config is not None and version in self.configs:
+      return config in self.configs[version]
+    return False
+
+
 class Experiment(enum.Enum):
   """Experiments which can be enabled/disabled on a per version basis.
 

tensorflow_datasets/core/utils/version_test.py

@@ -14,6 +14,8 @@
 # limitations under the License.
 
 """Tests for tensorflow_datasets.core.utils.version."""
+
+import pytest
 from tensorflow_datasets import testing
 from tensorflow_datasets.core.utils import version
 
@@ -148,5 +150,30 @@ def test_str_to_version():
   assert v1 == v0
 
 
+@pytest.mark.parametrize(
+    'blocked_version, blocked_config, expected',
+    [
+        ('1.2.3', None, True),
+        ('1.0.0', None, False),
+        ('1.2.3', 'non_existing_config', True),
+        ('1.0.0', 'config_1', True),
+        ('1.0.0', 'config_2', True),
+        ('1.0.0', 'config_1', True),
+        ('1.1.0', 'config_2', False),
+    ],
+)
+def test_is_blocked(blocked_version, blocked_config, expected):
+  blocked_versions = version.BlockedVersions(
+      versions={'1.2.3': None},
+      configs={
+          '1.0.0': {'config_1': 'blocked_config', 'config_2': None},
+          '1.1.0': {'config_1': 'blocked config in version 1.1.0'},
+      },
+  )
+  assert (
+      blocked_versions.is_blocked(blocked_version, blocked_config) == expected
+  )
+
+
 if __name__ == '__main__':
   testing.test_main()