Fix bug where the builder config was not retrieved in the right manner since the config name also includes the version

PiperOrigin-RevId: 694039995

Author: tomvdw

tensorflow_datasets/core/dataset_builder.py

@@ -21,6 +21,7 @@
 import collections
 from collections.abc import Iterable, Iterator, Mapping, Sequence
 import dataclasses
+import difflib
 import functools
 import inspect
 import json
@@ -1347,39 +1348,50 @@ def _create_builder_config(
   ) -> BuilderConfig | None:
     """Create and validate BuilderConfig object."""
     if builder_config is None:
-      builder_config = self.get_default_builder_config()
-      if builder_config is not None:
+      if builder_config := self.get_default_builder_config():
         logging.info(
             "No config specified, defaulting to config: %s/%s",
             self.name,
             builder_config.name,
         )
-    if not builder_config:
-      return None
+        return builder_config
+      else:
+        return None
+
     if isinstance(builder_config, str):
-      name = builder_config
-      builder_config = self.builder_configs.get(name)
-      if builder_config is None and version is not None:
-        builder_config = self.builder_configs.get(f"{name}:{version}")
-      if builder_config is None:
-        raise ValueError(
-            "BuilderConfig %s not found with version %s. Available: %s"
-            % (name, version, list(self.builder_configs.keys()))
+      config = self.get_builder_config(name=builder_config, version=version)
+      if config is not None:
+        return config
+      else:
+        close_matches = difflib.get_close_matches(
+            builder_config, self.builder_configs.keys(), n=10
         )
-    name = builder_config.name
-    if not name:
-      raise ValueError("BuilderConfig must have a name, got %s" % name)
-    is_custom = name not in self.builder_configs
-    if is_custom:
-      logging.warning("Using custom data configuration %s", name)
-    else:
-      if builder_config is not self.builder_configs[name]:
         raise ValueError(
-            "Cannot name a custom BuilderConfig the same as an available "
-            "BuilderConfig. Change the name. Available BuilderConfigs: %s"
-            % (list(self.builder_configs.keys()))
+            f"BuilderConfig {builder_config} not found with version {version}."
+            " Here are 10 BuilderConfigs whose name closely match:"
+            f" {close_matches}"
         )
-    return builder_config
+
+    if not isinstance(builder_config, BuilderConfig):
+      raise ValueError(
+          f"BuilderConfig {builder_config} is not a BuilderConfig or string"
+      )
+
+    cls_builder_config = self.get_builder_config(
+        name=builder_config.name, version=version
+    )
+    if cls_builder_config is None:
+      logging.warning("Using custom data configuration: %s", builder_config)
+      return builder_config
+    elif builder_config is not cls_builder_config:
+      raise ValueError(
+          "Cannot name a custom BuilderConfig the same as an available"
+          " BuilderConfig. Change the name.\n"
+          f"Requested: {builder_config}\n"
+          f"BuilderConfig in class: {cls_builder_config}"
+      )
+    else:
+      return builder_config
 
   @utils.classproperty
   @classmethod
@@ -2019,7 +2031,7 @@ def canonical_version_for_config(
 ) -> utils.Version:
   """Get the canonical version for the given config.
 
-  This allow to get the version without instanciating the class.
+  This allows getting the version without instantiating the class.
   The version can be stored either at the class or in the config object.
 
   Args:
@@ -2030,14 +2042,15 @@ def canonical_version_for_config(
   Returns:
     version: The extracted version.
   """
-  if instance_or_cls.BUILDER_CONFIGS and config is None:
+  if config and config.version:
+    return utils.Version(config.version)
+
+  if config is None and instance_or_cls.BUILDER_CONFIGS:
     raise ValueError(
         f"Cannot infer version on {instance_or_cls.name}. Unknown config."
     )
 
-  if config and config.version:
-    return utils.Version(config.version)
-  elif instance_or_cls.VERSION:
+  if instance_or_cls.VERSION:
     return utils.Version(instance_or_cls.VERSION)
   else:
     raise ValueError(

tensorflow_datasets/core/utils/file_utils.py

@@ -247,10 +247,14 @@ def get_data_dir_and_dataset_dir(
   if all_versions:
     logging.warning(
         (
-            'Found a different version of the requested dataset:\n'
-            '%s\n'
-            'Using %s instead.'
+            'Found a different version of the requested dataset'
+            ' (given_data_dir=%s,dataset=%s, config=%s, version=%s):\n'
+            '%s\nUsing %s instead.'
         ),
+        given_data_dir,
+        builder_name,
+        config_name,
+        version,
         '\n'.join(str(v) for v in sorted(all_versions)),
         dataset_dir,
     )

tensorflow_datasets/scripts/cli/build.py

@@ -16,13 +16,14 @@
 """`tfds build` command."""
 
 import argparse
+from collections.abc import Iterator
 import functools
 import importlib
 import itertools
 import json
 import multiprocessing
 import os
-from typing import Any, Dict, Iterator, Optional, Tuple, Type, Union
+from typing import Any, Type
 
 from absl import logging
 import tensorflow_datasets as tfds
@@ -132,7 +133,7 @@ def _build_datasets(args: argparse.Namespace) -> None:
 def _make_builders(
     args: argparse.Namespace,
     builder_cls: Type[tfds.core.DatasetBuilder],
-    builder_kwargs: Dict[str, Any],
+    builder_kwargs: dict[str, Any],
 ) -> Iterator[tfds.core.DatasetBuilder]:
   """Yields builders to generate.
 
@@ -185,7 +186,7 @@ def _get_builder_cls_and_kwargs(
     ds_to_build: str,
     *,
     has_imports: bool,
-) -> Tuple[Type[tfds.core.DatasetBuilder], Dict[str, Any]]:
+) -> tuple[Type[tfds.core.DatasetBuilder], dict[str, Any]]:
   """Infer the builder class to build and its kwargs.
 
   Args:
@@ -232,7 +233,7 @@ def _get_builder_cls_and_kwargs(
   return builder_cls, builder_kwargs
 
 
-def _search_script_path(ds_to_build: str) -> Optional[tfds.core.Path]:
+def _search_script_path(ds_to_build: str) -> tfds.core.Path | None:
   """Check whether the requested dataset match a file on disk."""
   # If the dataset file exists, use it. Valid values are:
   # * Empty string (use default directory)
@@ -274,7 +275,7 @@ def _search_script_path(ds_to_build: str) -> Optional[tfds.core.Path]:
 
 def _validate_script_path(
     path: tfds.core.Path,
-) -> Optional[tfds.core.Path]:
+) -> tfds.core.Path | None:
   """Validates and returns the `dataset_builder.py` generation script path."""
   if path.suffix != '.py':
     raise ValueError(f'Expected `.py` file. Invalid dataset path: {path}.')
@@ -292,17 +293,17 @@ def _make_builder(
     **builder_kwargs,
 ) -> tfds.core.DatasetBuilder:
   """Builder factory, eventually deleting pre-existing dataset."""
-  builder = builder_cls(**builder_kwargs)  # pytype: disable=not-instantiable
-  data_exists = builder.data_path.exists()
-  if fail_if_exists and data_exists:
+  builder = builder_cls(**builder_kwargs)
+  is_prepared = builder.is_prepared()
+  if fail_if_exists and is_prepared:
     raise RuntimeError(
         'The `fail_if_exists` flag was True and '
         f'the data already exists in {builder.data_path}'
     )
-  if overwrite and data_exists:
+  if overwrite and is_prepared:
     builder.data_path.rmtree()  # Delete pre-existing data
     # Re-create the builder with clean state
-    builder = builder_cls(**builder_kwargs)  # pytype: disable=not-instantiable
+    builder = builder_cls(**builder_kwargs)
   return builder
 
 
@@ -404,10 +405,10 @@ def _make_download_config(
 
 def _get_config_name(
     builder_cls: Type[tfds.core.DatasetBuilder],
-    config_kwarg: Optional[str],
-    config_name: Optional[str],
-    config_idx: Optional[int],
-) -> Optional[Union[str, tfds.core.BuilderConfig]]:
+    config_kwarg: str | None,
+    config_name: str | None,
+    config_idx: int | None,
+) -> str | tfds.core.BuilderConfig | None:
   """Extract the config name.
 
   Args: