Move _build_data_dir from dataset_builder.py to file_utils.py

PiperOrigin-RevId: 683968893

Author: fineguy

tensorflow_datasets/core/dataset_builder.py

@@ -295,13 +295,12 @@ def __init__(
     else:  # Use the code version (do not restore data)
       self.info.initialize_from_bucket()
     if self.BLOCKED_VERSIONS is not None:
-      config_name = self._builder_config.name if self._builder_config else None
       if is_blocked := self.BLOCKED_VERSIONS.is_blocked(
-          version=self._version, config=config_name
+          version=self._version, config=self.builder_config_name
       ):
         default_msg = (
             f"Dataset {self.name} is blocked at version {self._version} and"
-            f" config {config_name}."
+            f" config {self.builder_config_name}."
         )
         self.info.set_is_blocked(
             is_blocked.blocked_msg if is_blocked.blocked_msg else default_msg
@@ -1129,73 +1128,24 @@ def _should_cache_ds(self, split, shuffle_files, read_config) -> bool:
     # If the dataset satisfy all the right conditions, activate autocaching.
     return True
 
-  def _relative_data_dir(self, with_version: bool = True) -> str:
-    """Relative path of this dataset in data_dir."""
-    builder_data_dir = self.name
-    builder_config = self._builder_config
-    if builder_config:
-      builder_data_dir = os.path.join(builder_data_dir, builder_config.name)
-    if not with_version:
-      return builder_data_dir
-
-    version_data_dir = os.path.join(builder_data_dir, str(self._version))
-    return version_data_dir
-
-  def _build_data_dir(self, given_data_dir: Optional[str]):
+  def _build_data_dir(self, given_data_dir: str | None) -> tuple[str, str]:
     """Return the data directory for the current version.
 
     Args:
-      given_data_dir: `Optional[str]`, root `data_dir` passed as `__init__`
-        argument.
+      given_data_dir: Root `data_dir` passed as `__init__` argument.
 
     Returns:
-      data_dir_root: `str`, The root dir containing all datasets, downloads,...
-      data_dir: `str`, The version data_dir
-        (e.g. `<data_dir_root>/<ds_name>/<config>/<version>`)
+      data_dir: Root directory containing all datasets, downloads,...
+      dataset_dir: Dataset data directory (e.g.
+        `<data_dir>/<ds_name>/<config>/<version>`)
     """
-    builder_dir = self._relative_data_dir(with_version=False)
-    version_dir = self._relative_data_dir(with_version=True)
-
-    default_data_dir = file_utils.get_default_data_dir(
-        given_data_dir=given_data_dir
+    data_dir, dataset_dir = file_utils.get_data_dir_and_dataset_dir(
+        given_data_dir=given_data_dir,
+        builder_name=self.name,
+        config_name=self.builder_config_name,
+        version=self.version,
     )
-    all_data_dirs = file_utils.list_data_dirs(given_data_dir=given_data_dir)
-
-    all_versions = set()
-    requested_version_dirs = {}
-    for data_dir_root in all_data_dirs:
-      # List all existing versions
-      full_builder_dir = os.path.join(data_dir_root, builder_dir)
-      data_dir_versions = set(utils.version.list_all_versions(full_builder_dir))
-      # Check for existence of the requested version
-      if self.version in data_dir_versions:
-        requested_version_dirs[data_dir_root] = os.path.join(
-            data_dir_root, version_dir
-        )
-      all_versions.update(data_dir_versions)
-
-    if len(requested_version_dirs) > 1:
-      raise ValueError(
-          "Dataset was found in more than one directory: {}. Please resolve "
-          "the ambiguity by explicitly specifying `data_dir=`."
-          "".format(requested_version_dirs.values())
-      )
-    elif len(requested_version_dirs) == 1:  # The dataset is found once
-      return next(iter(requested_version_dirs.items()))
-
-    # No dataset found, use default directory
-    data_dir = os.path.join(default_data_dir, version_dir)
-    if all_versions:
-      logging.warning(
-          (
-              "Found a different version of the requested dataset:\n"
-              "%s\n"
-              "Using %s instead."
-          ),
-          "\n".join(str(v) for v in sorted(all_versions)),
-          data_dir,
-      )
-    return default_data_dir, data_dir
+    return os.fspath(data_dir), os.fspath(dataset_dir)
 
   def _log_download_done(self) -> None:
     msg = (
@@ -2032,13 +1982,9 @@ def _save_default_config_name(
     tmp_config_path.write_text(json.dumps(data))
 
 
-def load_default_config_name(
-    common_dir: epath.Path,
-) -> Optional[str]:
+def load_default_config_name(builder_dir: epath.Path) -> str | None:
   """Load `builder_cls` metadata (common to all builder configs)."""
-  config_path = (
-      epath.Path(common_dir) / f".config/{constants.METADATA_FILENAME}"
-  )
+  config_path = builder_dir / ".config" / constants.METADATA_FILENAME
   if not config_path.exists():
     return None
   data = json.loads(config_path.read_text())

tensorflow_datasets/core/dataset_builder_test.py

@@ -13,8 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Tests for tensorflow_datasets.core.dataset_builder."""
-
 from collections.abc import Iterator, Mapping, Sequence
 import dataclasses
 import functools
@@ -353,28 +351,34 @@ def test_build_data_dir(self):
     with testing.tmp_dir(self.get_temp_dir()) as tmp_dir:
       builder = DummyDatasetSharedGenerator(data_dir=tmp_dir)
       self.assertEqual(str(builder.info.version), "1.0.0")
-      builder_data_dir = os.path.join(tmp_dir, builder.name)
-      version_dir = os.path.join(builder_data_dir, "1.0.0")
+
+      builder_dir = file_utils.get_dataset_dir(
+          data_dir=tmp_dir, builder_name=builder.name
+      )
+      version_dir = builder_dir / "1.0.0"
 
       # The dataset folder contains multiple other versions
-      tf.io.gfile.makedirs(os.path.join(builder_data_dir, "14.0.0.invalid"))
-      tf.io.gfile.makedirs(os.path.join(builder_data_dir, "10.0.0"))
-      tf.io.gfile.makedirs(os.path.join(builder_data_dir, "9.0.0"))
-      tf.io.gfile.makedirs(os.path.join(builder_data_dir, "0.1.0"))
+      (builder_dir / "14.0.0.invalid").mkdir(parents=True, exist_ok=True)
+      (builder_dir / "10.0.0").mkdir(parents=True, exist_ok=True)
+      (builder_dir / "9.0.0").mkdir(parents=True, exist_ok=True)
+      (builder_dir / "0.1.0").mkdir(parents=True, exist_ok=True)
 
       # The builder's version dir is chosen
-      self.assertEqual(builder._build_data_dir(tmp_dir)[1], version_dir)
+      self.assertEqual(builder.data_path, version_dir)
 
   def test_get_data_dir_with_config(self):
     with testing.tmp_dir(self.get_temp_dir()) as tmp_dir:
       config_name = "plus1"
       builder = DummyDatasetWithConfigs(config=config_name, data_dir=tmp_dir)
 
-      builder_data_dir = os.path.join(tmp_dir, builder.name, config_name)
-      version_data_dir = os.path.join(builder_data_dir, "0.0.1")
+      version_dir = file_utils.get_dataset_dir(
+          data_dir=tmp_dir,
+          builder_name=builder.name,
+          config_name=config_name,
+          version="0.0.1",
+      )
 
-      tf.io.gfile.makedirs(version_data_dir)
-      self.assertEqual(builder._build_data_dir(tmp_dir)[1], version_data_dir)
+      self.assertEqual(builder.data_path, version_dir)
 
   def test_config_construction(self):
     with testing.tmp_dir(self.get_temp_dir()) as tmp_dir:
@@ -775,12 +779,11 @@ class DatasetBuilderMultiDirTest(testing.TestCase):
   def setUpClass(cls):
     super(DatasetBuilderMultiDirTest, cls).setUpClass()
     cls.builder = DummyDatasetSharedGenerator()
-    cls.version_dir = os.path.normpath(cls.builder.info.full_name)
 
   def setUp(self):
     super(DatasetBuilderMultiDirTest, self).setUp()
     # Sanity check to make sure that no dir is registered
-    self.assertEmpty(file_utils._registered_data_dir)
+    file_utils.clear_registered_data_dirs()
     # Create a new temp dir
     self.other_data_dir = os.path.join(self.get_temp_dir(), "other_dir")
     # Overwrite the default data_dir (as files get created)
@@ -801,102 +804,6 @@ def tearDown(self):
     # Restore the orgininal data dir
     constants.DATA_DIR = self._original_data_dir
 
-  def assertBuildDataDir(self, build_data_dir_out, root_dir):
-    data_dir_root, data_dir = build_data_dir_out
-    self.assertEqual(data_dir_root, root_dir)
-    self.assertEqual(data_dir, os.path.join(root_dir, self.version_dir))
-
-  def test_default(self):
-    # No data_dir is passed
-    # -> use default path is used.
-    self.assertBuildDataDir(
-        self.builder._build_data_dir(None), self.default_data_dir
-    )
-
-  def test_explicitly_passed(self):
-    # When a dir is explicitly passed, use it.
-    self.assertBuildDataDir(
-        self.builder._build_data_dir(self.other_data_dir), self.other_data_dir
-    )
-
-  def test_default_multi_dir(self):
-    # No data_dir is passed
-    # Multiple data_dirs are registered
-    # -> use default path
-    file_utils.add_data_dir(self.other_data_dir)
-    self.assertBuildDataDir(
-        self.builder._build_data_dir(None), self.default_data_dir
-    )
-
-  def test_default_multi_dir_old_version_exists(self):
-    # No data_dir is passed
-    # Multiple data_dirs are registered
-    # Data dir contains old versions
-    # -> use default path
-    file_utils.add_data_dir(self.other_data_dir)
-    tf.io.gfile.makedirs(
-        os.path.join(
-            self.other_data_dir, "dummy_dataset_shared_generator", "0.1.0"
-        )
-    )
-    tf.io.gfile.makedirs(
-        os.path.join(
-            self.other_data_dir, "dummy_dataset_shared_generator", "0.2.0"
-        )
-    )
-    self.assertBuildDataDir(
-        self.builder._build_data_dir(None), self.default_data_dir
-    )
-
-  def test_default_multi_dir_version_exists(self):
-    # No data_dir is passed
-    # Multiple data_dirs are registered
-    # Data found
-    # -> Re-load existing data
-    file_utils.add_data_dir(self.other_data_dir)
-    tf.io.gfile.makedirs(
-        os.path.join(
-            self.other_data_dir, "dummy_dataset_shared_generator", "1.0.0"
-        )
-    )
-    self.assertBuildDataDir(
-        self.builder._build_data_dir(None), self.other_data_dir
-    )
-
-  def test_default_multi_dir_duplicate(self):
-    # If two data dirs contains the dataset, raise an error...
-    file_utils.add_data_dir(self.other_data_dir)
-    tf.io.gfile.makedirs(
-        os.path.join(
-            self.default_data_dir, "dummy_dataset_shared_generator", "1.0.0"
-        )
-    )
-    tf.io.gfile.makedirs(
-        os.path.join(
-            self.other_data_dir, "dummy_dataset_shared_generator", "1.0.0"
-        )
-    )
-    with self.assertRaisesRegex(ValueError, "found in more than one directory"):
-      self.builder._build_data_dir(None)
-
-  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)
-    tf.io.gfile.makedirs(
-        os.path.join(
-            self.default_data_dir, "dummy_dataset_shared_generator", "1.0.0"
-        )
-    )
-    tf.io.gfile.makedirs(
-        os.path.join(
-            self.other_data_dir, "dummy_dataset_shared_generator", "1.0.0"
-        )
-    )
-    self.assertBuildDataDir(
-        self.builder._build_data_dir(self.other_data_dir), self.other_data_dir
-    )
-
   def test_load_data_dir(self):
     """Ensure that `tfds.load` also supports multiple data_dir."""
     file_utils.add_data_dir(self.other_data_dir)