Add overwrite option to utils.incomplete_dir().

PiperOrigin-RevId: 686088351

Author: fineguy

tensorflow_datasets/core/dataset_builder.py

@@ -627,15 +627,16 @@ def download_and_prepare(
       self._update_dataset_info()
       return
 
-    if data_exists and download_config.download_mode == REUSE_DATASET_IF_EXISTS:
-      logging.info("Reusing dataset %s (%s)", self.name, self.data_dir)
-      return
-    elif data_exists and download_config.download_mode == REUSE_CACHE_IF_EXISTS:
-      logging.info(
-          "Deleting pre-existing dataset %s (%s)", self.name, self.data_dir
-      )
-      data_path.rmtree()  # Delete pre-existing data.
-      data_exists = data_path.exists()
+    if data_exists:
+      if download_config.download_mode.overwrite_dataset:
+        logging.info(
+            "Deleting pre-existing dataset %s (%s)", self.name, self.data_dir
+        )
+        data_path.rmtree()  # Delete pre-existing data.
+        data_exists = data_path.exists()
+      else:
+        logging.info("Reusing dataset %s (%s)", self.name, self.data_dir)
+        return
 
     if self.version.tfds_version_to_prepare:
       available_to_prepare = ", ".join(
@@ -734,7 +735,9 @@ def download_and_prepare(
 
     # Create a tmp dir and rename to self.data_dir on successful exit.
     with utils.incomplete_dir(
-        dirname=self.data_dir, permissions=permissions
+        dirname=self.data_dir,
+        permissions=permissions,
+        overwrite=download_config.download_mode.overwrite_dataset,
     ) as tmp_data_dir:
       # Temporarily assign _data_dir to tmp_data_dir to avoid having to forward
       # it to every sub function.
@@ -1297,8 +1300,8 @@ def _make_download_manager(
         manual_dir=manual_dir,
         url_infos=self.url_infos,
         manual_dir_instructions=self.MANUAL_DOWNLOAD_INSTRUCTIONS,
-        force_download=(download_config.download_mode == FORCE_REDOWNLOAD),
-        force_extraction=(download_config.download_mode == FORCE_REDOWNLOAD),
+        force_download=download_config.download_mode.force_download,
+        force_extraction=download_config.download_mode.force_download,
         force_checksums_validation=download_config.force_checksums_validation,
         register_checksums=download_config.register_checksums,
         register_checksums_path=register_checksums_path,

tensorflow_datasets/core/download/util.py

@@ -16,13 +16,14 @@
 """Utils functions."""
 
 import enum
+from etils import epy
 
 
 class DownloadError(Exception):
   pass
 
 
-class GenerateMode(enum.Enum):
+class GenerateMode(epy.StrEnum):
   """`Enum` for how to treat pre-existing downloads and data.
 
   The default mode is `REUSE_DATASET_IF_EXISTS`, which will reuse both
@@ -46,10 +47,21 @@ class GenerateMode(enum.Enum):
   UPDATE_DATASET_INFO will fail if the data has never been prepared.
   """
 
-  REUSE_DATASET_IF_EXISTS = 'reuse_dataset_if_exists'
-  UPDATE_DATASET_INFO = 'update_dataset_info'
-  REUSE_CACHE_IF_EXISTS = 'reuse_cache_if_exists'
-  FORCE_REDOWNLOAD = 'force_redownload'
+  REUSE_DATASET_IF_EXISTS = enum.auto()
+  UPDATE_DATASET_INFO = enum.auto()
+  REUSE_CACHE_IF_EXISTS = enum.auto()
+  FORCE_REDOWNLOAD = enum.auto()
+
+  @property
+  def force_download(self) -> bool:
+    return self == GenerateMode.FORCE_REDOWNLOAD
+
+  @property
+  def overwrite_dataset(self) -> bool:
+    return self in [
+        GenerateMode.REUSE_CACHE_IF_EXISTS,
+        GenerateMode.FORCE_REDOWNLOAD,
+    ]
 
 
 class ComputeStatsMode(enum.Enum):

tensorflow_datasets/core/utils/file_utils.py

@@ -76,7 +76,7 @@ def as_path(path: PathLike) -> Path:
       """
   )
   logging.warning(msg)
-  return epath.Path(path)
+  return Path(path)
 
 
 def add_data_dir(data_dir: PathLike) -> None:
@@ -94,7 +94,7 @@ def add_data_dir(data_dir: PathLike) -> None:
   Args:
     data_dir: New data_dir to register.
   """
-  _REGISTERED_DATA_DIRS.add(epath.Path(data_dir))
+  _REGISTERED_DATA_DIRS.add(Path(data_dir))
 
 
 def clear_registered_data_dirs() -> None:
@@ -107,21 +107,25 @@ def _get_incomplete_dir(dir_name: str) -> str:
   random_suffix = ''.join(
       random.choice(string.ascii_uppercase + string.digits) for _ in range(6)
   )
-  dir_name = epath.Path(dir_name)
+  dir_name = Path(dir_name)
   return f'{dir_name.parent}/{constants.INCOMPLETE_PREFIX}{random_suffix}_{dir_name.name}/'
 
 
 @contextlib.contextmanager
 def incomplete_dir(
-    dirname: epath.PathLike, permissions: Permissions | None = None
+    dirname: PathLike,
+    permissions: Permissions | None = None,
+    overwrite: bool = False,
 ) -> Iterator[str]:
   """Create temporary dir for dirname and rename on exit."""
   dirname = os.fspath(dirname)
   tmp_dir = _get_incomplete_dir(dirname)
-  tmp_path = epath.Path(tmp_dir)
+  tmp_path = Path(tmp_dir)
   tmp_path.mkdir(parents=True, exist_ok=True)
   try:
     yield tmp_dir
+    if overwrite:
+      Path(dirname).rmtree(missing_ok=True)
     tmp_path.rename(dirname)
   finally:
     if tmp_path.exists():

tensorflow_datasets/scripts/cli/cli_utils.py

@@ -315,6 +315,8 @@ def download_and_prepare(
 
   if not download_config:
     download_config = download.DownloadConfig()
+  if overwrite and not download_config.download_mode.overwrite_dataset:
+    download_config.download_mode = download.GenerateMode.REUSE_CACHE_IF_EXISTS
 
   # Add Apache Beam options to download config
   try: