Add VersionOrStr type annotations.

PiperOrigin-RevId: 693252714

Author: fineguy

tensorflow_datasets/core/dataset_builder.py

@@ -69,7 +69,6 @@
 ListOrTreeOrElem = type_utils.ListOrTreeOrElem
 Tree = type_utils.Tree
 TreeDict = type_utils.TreeDict
-VersionOrStr = utils.Version | str
 
 FORCE_REDOWNLOAD = download.GenerateMode.FORCE_REDOWNLOAD
 REUSE_CACHE_IF_EXISTS = download.GenerateMode.REUSE_CACHE_IF_EXISTS
@@ -107,9 +106,9 @@ class BuilderConfig:
   # * Kwargs-only (https://bugs.python.org/issue33129)
 
   name: str
-  version: VersionOrStr | None = None
+  version: utils.VersionOrStr | None = None
   release_notes: dict[str, str] | None = None
-  supported_versions: list[VersionOrStr] = dataclasses.field(
+  supported_versions: list[utils.VersionOrStr] = dataclasses.field(
       default_factory=list
   )
   description: str | None = None

tensorflow_datasets/core/dataset_builders/croissant_builder.py

@@ -55,7 +55,7 @@
 from tensorflow_datasets.core.utils import conversion_utils
 from tensorflow_datasets.core.utils import croissant_utils
 from tensorflow_datasets.core.utils import type_utils
-from tensorflow_datasets.core.utils import version as version_utils
+from tensorflow_datasets.core.utils import version as version_lib
 from tensorflow_datasets.core.utils.lazy_imports_utils import apache_beam as beam
 from tensorflow_datasets.core.utils.lazy_imports_utils import mlcroissant as mlc
 from tensorflow_datasets.core.utils.lazy_imports_utils import pandas as pd
@@ -165,7 +165,7 @@ def __init__(
       int_dtype: type_utils.TfdsDType | None = np.int64,
       float_dtype: type_utils.TfdsDType | None = np.float32,
       mapping: Mapping[str, epath.PathLike] | None = None,
-      overwrite_version: str | None = None,
+      overwrite_version: version_lib.VersionOrStr | None = None,
       filters: Mapping[str, Any] | None = None,
       **kwargs: Any,
   ):
@@ -203,7 +203,7 @@ def __init__(
     # In TFDS, version is a mandatory attribute, while in Croissant it is only a
     # recommended attribute. If the version is unspecified in Croissant, we set
     # it to `1.0.0` in TFDS.
-    self.VERSION = version_utils.Version(  # pylint: disable=invalid-name
+    self.VERSION = version_lib.Version(  # pylint: disable=invalid-name
         overwrite_version or self.dataset.metadata.version or '1.0.0'
     )
     self.RELEASE_NOTES = {}  # pylint: disable=invalid-name

tensorflow_datasets/core/utils/__init__.py

@@ -88,3 +88,4 @@
 from tensorflow_datasets.core.utils.version import Experiment
 from tensorflow_datasets.core.utils.version import IsBlocked
 from tensorflow_datasets.core.utils.version import Version
+from tensorflow_datasets.core.utils.version import VersionOrStr

tensorflow_datasets/core/utils/version.py

@@ -20,6 +20,7 @@
 import dataclasses
 import enum
 import re
+from typing import Union
 
 from etils import epath
 
@@ -39,6 +40,7 @@ class DatasetVariantBlockedError(ValueError):
 # 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]
+VersionOrStr = Union["Version", str]
 
 
 @dataclasses.dataclass(frozen=True)
@@ -73,7 +75,7 @@ class BlockedVersions:
   configs: dict[str, BlockedWithMsg] = dataclasses.field(default_factory=dict)
 
   def is_blocked(
-      self, version: str | Version, config: str | None = None
+      self, version: VersionOrStr, config: str | None = None
   ) -> IsBlocked:
     """Checks whether a version or config is blocked.
 
@@ -131,7 +133,7 @@ class Version:
 
   def __init__(
       self,
-      version: Version | str,
+      version: VersionOrStr,
       experiments=None,
       tfds_version_to_prepare=None,
   ):
@@ -227,7 +229,7 @@ def match(self, other_version) -> bool:
     )
 
   @classmethod
-  def is_valid(cls, version: Version | str | None) -> bool:
+  def is_valid(cls, version: VersionOrStr | None) -> bool:
     """Returns True if the version can be parsed."""
     if isinstance(version, Version):
       return True

tensorflow_datasets/scripts/cli/croissant.py

@@ -38,6 +38,7 @@
 from tensorflow_datasets.core import file_adapters
 from tensorflow_datasets.core.dataset_builders import croissant_builder
 from tensorflow_datasets.core.utils import croissant_utils
+from tensorflow_datasets.core.utils import version as version_lib
 from tensorflow_datasets.scripts.cli import cli_utils
 
 
@@ -112,6 +113,15 @@ def record_set_ids(self) -> list[str]:
         self.dataset.metadata
     )
 
+  @functools.cached_property
+  def version(self) -> version_lib.Version:
+    # In TFDS, version is a mandatory attribute, while in Croissant it is only a
+    # recommended attribute. If the version is unspecified in Croissant, we set
+    # it to `1.0.0` in TFDS.
+    return version_lib.Version(
+        self.overwrite_version or self.dataset.metadata.version or '1.0.0'
+    )
+
 
 def register_subparser(parsers: argparse._SubParsersAction):
   """Add subparser for `convert_format` command."""
@@ -146,7 +156,7 @@ def prepare_croissant_builder(
       file_format=args.file_format,
       data_dir=args.data_dir,
       mapping=args.mapping_json,
-      overwrite_version=args.overwrite_version,
+      overwrite_version=args.version,
   )
   cli_utils.download_and_prepare(
       builder=builder,