Fix validating checksums path.

PiperOrigin-RevId: 670307758

Author: fineguy

tensorflow_datasets/core/dataset_builder.py

@@ -1280,7 +1280,11 @@ def _make_download_manager(
     if download_config.register_checksums:
       # Note: Error will be raised here if user try to record checksums
       # from a `zipapp`
-      register_checksums_path = utils.to_write_path(self._checksums_path)
+      try:
+        register_checksums_path = utils.to_write_path(self._checksums_path)
+        download.validate_checksums_path(register_checksums_path)
+      except Exception:  # pylint: disable=broad-except
+        raise
     else:
       register_checksums_path = None
 

tensorflow_datasets/core/download/__init__.py

@@ -16,6 +16,7 @@
 """`tfds.download.DownloadManager` API."""
 
 from tensorflow_datasets.core.download.checksums import add_checksums_dir
+from tensorflow_datasets.core.download.checksums import validate_checksums_path
 from tensorflow_datasets.core.download.download_manager import DownloadConfig
 from tensorflow_datasets.core.download.download_manager import DownloadManager
 from tensorflow_datasets.core.download.extractor import iter_archive

tensorflow_datasets/core/download/checksums.py

@@ -216,6 +216,7 @@ def save_url_infos(
       f'{url_info.filename or ""}\n'
       for url, url_info in sorted(new_data.items())
   ]
+  path.parent.mkdir(parents=True, exist_ok=True)
   path.write_text(''.join(lines), encoding='UTF-8')
 
 
@@ -227,3 +228,19 @@ def _filenames_equal(
   return all(
       l.filename == r.filename for _, (l, r) in utils.zip_dict(left, right)
   )
+
+
+def validate_checksums_path(checksums_path: epath.PathLike):
+  """Validates the checksums path.
+
+  This function creates the file if it doesn't exist, and writes to it to make
+  sure the user has write access before downloading any files.
+
+  Args:
+    checksums_path: Path to the checksums file.
+  """
+  checksums_path = epath.Path(checksums_path)
+  if not checksums_path.exists():
+    checksums_path.touch()
+  else:
+    checksums_path.write_text(checksums_path.read_text())

tensorflow_datasets/core/download/download_manager.py

@@ -239,14 +239,6 @@ def __init__(
             ' set.'
         )
       register_checksums_path = epath.Path(register_checksums_path)
-      if not register_checksums_path.exists():
-        # Create the file here to make sure user has write access before
-        # starting downloads.
-        register_checksums_path.touch()
-      else:
-        # Make sure the user has write access before downloading any files.
-        # (e.g. TFDS installed by admin)
-        register_checksums_path.write_text(register_checksums_path.read_text())
 
     download_dir = epath.Path(download_dir).expanduser()
     if extract_dir: