Move file_utils.get_default_data_dir() to constants.get_default_data_dir().

PiperOrigin-RevId: 796869816

Author: fineguy

tensorflow_datasets/core/constants.py

@@ -28,10 +28,17 @@
 
 # Directory where to store processed datasets.
 # If modifying this, should also update `scripts/cli/build.py` `--data_dir`
-DATA_DIR: Final[str] = os.environ.get(
-    'TFDS_DATA_DIR',
-    os.path.join(os.path.expanduser('~'), 'tensorflow_datasets'),
-)
+
+
+def get_default_data_dir() -> str:
+  """Returns the TFDS default data directory."""
+  return os.environ.get(
+      'TFDS_DATA_DIR',
+      os.path.join(os.path.expanduser('~'), 'tensorflow_datasets'),
+  )
+
+
+DATA_DIR: Final[str] = get_default_data_dir()
 
 # Prefix of files / directories which aren't finished downloading / extracting.
 INCOMPLETE_PREFIX = 'incomplete.'

tensorflow_datasets/core/dataset_builder_test.py

@@ -27,7 +27,6 @@
 import pytest
 import tensorflow as tf
 from tensorflow_datasets import testing
-from tensorflow_datasets.core import constants
 from tensorflow_datasets.core import dataset_builder
 from tensorflow_datasets.core import dataset_info
 from tensorflow_datasets.core import dataset_utils
@@ -831,32 +830,19 @@ class DatasetBuilderMultiDirTest(testing.TestCase):
 
   @classmethod
   def setUpClass(cls):
-    super(DatasetBuilderMultiDirTest, cls).setUpClass()
+    super().setUpClass()
     cls.builder = DummyDatasetSharedGenerator()
 
   def setUp(self):
-    super(DatasetBuilderMultiDirTest, self).setUp()
+    super().setUp()
     # Sanity check to make sure that no dir is registered
     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)
-
-    self._original_data_dir = constants.DATA_DIR
-    constants.DATA_DIR = os.path.join(self.get_temp_dir(), "default_dir")
-    self.default_data_dir = constants.DATA_DIR
-
-  def tearDown(self):
-    super(DatasetBuilderMultiDirTest, self).tearDown()
-    # Restore to the default `_registered_data_dir`
-    file_utils._registered_data_dir = set()
-    # Clear-up existing dirs
-    if tf.io.gfile.exists(self.other_data_dir):
-      tf.io.gfile.rmtree(self.other_data_dir)
-    if tf.io.gfile.exists(self.default_data_dir):
-      tf.io.gfile.rmtree(self.default_data_dir)
-    # Restore the orgininal data dir
-    constants.DATA_DIR = self._original_data_dir
+    default_data_dir = self.enter_context(testing.mock_default_data_dir())
+    # Create a new temp dir
+    self.other_data_dir = os.path.join(
+        os.path.dirname(default_data_dir), "other_dir"
+    )
 
   def test_load_data_dir(self):
     """Ensure that `tfds.load` also supports multiple data_dir."""

tensorflow_datasets/core/dataset_builders/adhoc_builder.py

@@ -85,6 +85,7 @@
 
 from absl import logging
 from etils import epath
+from tensorflow_datasets.core import constants
 from tensorflow_datasets.core import dataset_builder
 from tensorflow_datasets.core import dataset_info
 from tensorflow_datasets.core import dataset_utils

tensorflow_datasets/core/utils/file_utils.py

@@ -156,23 +156,11 @@ def list_data_dirs(
     else:
       return [given_data_dir]
   else:
-    default_data_dir = get_default_data_dir(given_data_dir=given_data_dir)
+    default_data_dir = Path(constants.DATA_DIR)
     all_data_dirs = _REGISTERED_DATA_DIRS | {default_data_dir}
     return sorted(d.expanduser() for d in all_data_dirs)
 
 
-def get_default_data_dir(given_data_dir: epath.PathLike | None = None) -> Path:
-  """Returns the default data_dir."""
-  if given_data_dir:
-    data_dir = os.path.expanduser(given_data_dir)
-  elif 'TFDS_DATA_DIR' in os.environ:
-    data_dir = os.environ['TFDS_DATA_DIR']
-  else:
-    data_dir = constants.DATA_DIR
-
-  return Path(data_dir)
-
-
 def get_dataset_dir(
     data_dir: epath.PathLike,
     builder_name: str,
@@ -189,11 +177,11 @@ def get_dataset_dir(
 
 
 def get_data_dir_and_dataset_dir(
-    given_data_dir: epath.PathLike | None,
+    given_data_dir: PathLike | None,
     builder_name: str,
     config_name: str | None,
     version: version_lib.Version | str | None,
-) -> tuple[epath.Path, epath.Path]:
+) -> tuple[Path, Path]:
   """Returns the data and dataset directories for the given dataset.
 
   Args:
@@ -249,7 +237,7 @@ def get_data_dir_and_dataset_dir(
     return next(iter(dataset_dir_by_data_dir.items()))
 
   # No dataset found, use default directory
-  default_data_dir = get_default_data_dir()
+  default_data_dir = Path(constants.DATA_DIR)
   dataset_dir = get_dataset_dir(
       data_dir=default_data_dir,
       builder_name=builder_name,

tensorflow_datasets/core/utils/file_utils_test.py

@@ -33,7 +33,7 @@
 
 
 def test_default_data_dir():
-  data_dir = file_utils.get_default_data_dir(given_data_dir=None)
+  data_dir = constants.get_default_data_dir()
   assert data_dir
 
 

tensorflow_datasets/scripts/cli/cli_utils.py

@@ -197,7 +197,7 @@ class PathOptions:
   """
 
   data_dir: epath.Path = simple_parsing.field(
-      default=epath.Path(constants.DATA_DIR)
+      default=epath.Path(constants.get_default_data_dir())
   )
   download_dir: epath.Path | None = None
   extract_dir: epath.Path | None = None

tensorflow_datasets/testing/__init__.py

@@ -50,6 +50,7 @@
   from tensorflow_datasets.testing.test_utils import enable_gcs_access
   from tensorflow_datasets.testing.test_utils import fake_examples_dir
   from tensorflow_datasets.testing.test_utils import make_tmp_dir
+  from tensorflow_datasets.testing.test_utils import mock_default_data_dir
   from tensorflow_datasets.testing.test_utils import mock_kaggle_api
   from tensorflow_datasets.testing.test_utils import MockFs
   from tensorflow_datasets.testing.test_utils import rm_tmp_dir
@@ -84,6 +85,7 @@
     # TODO(afrozm): rm from here and add as methods to TestCase
     "make_tmp_dir": "tensorflow_datasets.testing.test_utils",
     "mock_data": "tensorflow_datasets.testing.mocking",
+    "mock_default_data_dir": "tensorflow_datasets.testing.test_utils",
     "mock_kaggle_api": "tensorflow_datasets.testing.test_utils",
     "MockFs": "tensorflow_datasets.testing.test_utils",
     "MockPolicy": "tensorflow_datasets.testing.mocking",

tensorflow_datasets/testing/test_utils.py

@@ -33,6 +33,7 @@
 from etils import epath
 from etils import epy
 import numpy as np
+from tensorflow_datasets.core import constants
 from tensorflow_datasets.core import dataset_builder
 from tensorflow_datasets.core import dataset_collection_builder
 from tensorflow_datasets.core import dataset_info
@@ -884,3 +885,13 @@ def dummy_croissant_file(
     croissant_file.write_text(json.dumps(dummy_metadata.to_json(), indent=2))
 
     yield croissant_file
+
+
+@contextlib.contextmanager
+def mock_default_data_dir() -> Iterator[str]:
+  """Mocks the `constants.DATA_DIR`."""
+  with tempfile.TemporaryDirectory() as tempdir:
+    tmp_data_dir = os.path.join(tempdir, 'default_dir')
+    os.makedirs(tmp_data_dir)
+    constants.DATA_DIR = tmp_data_dir
+    yield tmp_data_dir