Reuse pytest fixture for setting TFDS_DATA_DIR.

PiperOrigin-RevId: 792731989

Author: fineguy

tensorflow_datasets/conftest.py

@@ -20,16 +20,21 @@
 
 See: https://docs.pytest.org/en/latest/writing_plugins.html
 """
+
 from __future__ import annotations
 
 import builtins
 import importlib
+import os
+import pathlib
 import sys
 import typing
 from typing import Iterator, Type
 
+from etils import epath
 import pytest
 from tensorflow_datasets import setup_teardown
+from tensorflow_datasets.core import constants
 
 if typing.TYPE_CHECKING:
   from tensorflow_datasets import testing
@@ -52,11 +57,9 @@ def disable_community_datasets():
   # visibility isn't automatically set.
   from tensorflow_datasets.core import visibility  # pylint: disable=g-import-not-at-top
 
-  visibility.set_availables(
-      [
-          visibility.DatasetType.TFDS_PUBLIC,
-      ]
-  )
+  visibility.set_availables([
+      visibility.DatasetType.TFDS_PUBLIC,
+  ])
 
 
 @pytest.fixture(scope='session', autouse=True)
@@ -144,3 +147,13 @@ def dummy_dataset(
   return _make_dataset(tmp_path_factory, testing.DummyDataset)
 
 
+@pytest.fixture(name='default_data_dir')
+def mock_default_data_dir(
+    monkeypatch: pytest.MonkeyPatch, tmp_path: pathlib.Path
+) -> epath.Path:
+  """Sets the default data dir to a temp dir."""
+  default_data_dir = epath.Path(tmp_path) / 'default_data_dir'
+  monkeypatch.setattr(constants, 'DATA_DIR', os.fspath(default_data_dir))
+  return default_data_dir
+
+

tensorflow_datasets/core/constants.py

@@ -18,6 +18,7 @@
 # IMPORTANT: when changing values here, update docstrings.
 
 import os
+from typing import Final
 
 # Directory in which datasets are declared within TFDS sources.
 DATASETS_TFDS_SRC_DIR = 'datasets'
@@ -27,7 +28,7 @@
 
 # Directory where to store processed datasets.
 # If modifying this, should also update `scripts/cli/build.py` `--data_dir`
-DATA_DIR = os.environ.get(
+DATA_DIR: Final[str] = os.environ.get(
     'TFDS_DATA_DIR',
     os.path.join(os.path.expanduser('~'), 'tensorflow_datasets'),
 )

tensorflow_datasets/core/utils/file_utils_test.py

@@ -62,16 +62,8 @@ def _assert_data_dir(
   assert data_dir == expected_data_dir
 
 
-@pytest.fixture(name='default_data_dir')
-def mock_default_data_dir(monkeypatch, tmp_path):
-  """Sets the default data dir to a temp dir."""
-  default_data_dir = tmp_path / 'default_data_dir'
-  monkeypatch.setattr(constants, 'DATA_DIR', default_data_dir)
-  return default_data_dir
-
-
 @pytest.fixture(name='other_data_dir')
-def mock_other_data_dir(default_data_dir):
+def mock_other_data_dir(default_data_dir: epath.Path):
   """Adds another data dir to the registered data dirs."""
   other_data_dir = default_data_dir.parent / 'other_data_dir'
   file_utils.add_data_dir(other_data_dir)