Add option to download_and_prepare to set the permissions of newly created folders and files

Moved `incomplete_dir` to file_utils because I think it fits there better.

PiperOrigin-RevId: 652453020

Author: tomvdw

tensorflow_datasets/core/__init__.py

@@ -49,6 +49,7 @@
 from tensorflow_datasets.core.utils.benchmark import BenchmarkResult
 from tensorflow_datasets.core.utils.file_utils import add_data_dir
 from tensorflow_datasets.core.utils.file_utils import as_path
+from tensorflow_datasets.core.utils.file_utils import Permissions
 from tensorflow_datasets.core.writer import ExampleWriter
 
 
@@ -79,6 +80,7 @@ def benchmark(*args, **kwargs):
     "Metadata",
     "MetadataDict",
     "Path",
+    "Permissions",
     "ReadInstruction",
     "SequentialWriter",
     "ShardedFileTemplate",

tensorflow_datasets/core/dataset_builder.py

@@ -561,9 +561,10 @@ def is_prepared(self) -> bool:
   def download_and_prepare(
       self,
       *,
-      download_dir: Optional[epath.PathLike] = None,
-      download_config: Optional[download.DownloadConfig] = None,
-      file_format: Optional[Union[str, file_adapters.FileFormat]] = None,
+      download_dir: epath.PathLike | None = None,
+      download_config: download.DownloadConfig | None = None,
+      file_format: str | file_adapters.FileFormat | None = None,
+      permissions: file_utils.Permissions | None = None,
   ) -> None:
     """Downloads and prepares dataset for reading.
 
@@ -574,6 +575,8 @@ def download_and_prepare(
         downloading and preparing dataset.
       file_format: optional `str` or `file_adapters.FileFormat`, format of the
         record files in which the dataset will be written.
+      permissions: optional permissions to set on the generated folder and
+        files.
 
     Raises:
       IOError: if there is not enough disk space available.
@@ -694,7 +697,9 @@ def download_and_prepare(
       self.info.set_file_format(file_format, override=True)
 
     # Create a tmp dir and rename to self.data_dir on successful exit.
-    with utils.incomplete_dir(self.data_dir) as tmp_data_dir:
+    with utils.incomplete_dir(
+        dirname=self.data_dir, permissions=permissions
+    ) as tmp_data_dir:
       # Temporarily assign _data_dir to tmp_data_dir to avoid having to forward
       # it to every sub function.
       with utils.temporary_assignment(self, "_data_dir", tmp_data_dir):

tensorflow_datasets/core/utils/__init__.py

@@ -17,6 +17,7 @@
 
 from tensorflow_datasets.core.units import Size
 from tensorflow_datasets.core.utils import docs
+from tensorflow_datasets.core.utils.file_utils import incomplete_dir
 from tensorflow_datasets.core.utils.gcs_utils import gcs_path
 from tensorflow_datasets.core.utils.image_utils import apply_colormap
 from tensorflow_datasets.core.utils.image_utils import create_thumbnail
@@ -40,7 +41,6 @@
 from tensorflow_datasets.core.utils.py_utils import get_class_path
 from tensorflow_datasets.core.utils.py_utils import get_class_url
 from tensorflow_datasets.core.utils.py_utils import has_sufficient_disk_space
-from tensorflow_datasets.core.utils.py_utils import incomplete_dir
 from tensorflow_datasets.core.utils.py_utils import incomplete_file
 from tensorflow_datasets.core.utils.py_utils import indent
 from tensorflow_datasets.core.utils.py_utils import is_incomplete_file

tensorflow_datasets/core/utils/file_utils.py

@@ -19,9 +19,13 @@
 
 import collections
 from collections.abc import Iterator, Sequence
+import contextlib
+import dataclasses
 import functools
 import os
+import random
 import re
+import string
 import time
 
 from absl import logging
@@ -43,6 +47,15 @@
 _GLOB_CHARS = ['*', '?', '[']
 
 
+@dataclasses.dataclass(frozen=True)
+class Permissions:
+  """Permissions for a file or directory."""
+
+  owner: str | None = None
+  group: str | None = None
+  mode: int | None = None
+
+
 @docs.deprecated
 def as_path(path: PathLike) -> Path:
   """DEPRECATED. Please use `from etils import epath` with `epath.Path()`."""
@@ -87,6 +100,32 @@ def add_data_dir(data_dir):
   _registered_data_dir.add(data_dir)
 
 
+def _get_incomplete_dir(dir_name: str) -> str:
+  """Returns a temporary dir name based on `dir_name`."""
+  random_suffix = ''.join(
+      random.choice(string.ascii_uppercase + string.digits) for _ in range(6)
+  )
+  dir_name = epath.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
+) -> 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.mkdir(parents=True, exist_ok=True)
+  try:
+    yield tmp_dir
+    tmp_path.rename(dirname)
+  finally:
+    if tmp_path.exists():
+      tmp_path.rmtree()
+
+
 def list_data_dirs(
     given_data_dir: ListOrElem[PathLike] | None = None,
     dataset: str | None = None,

tensorflow_datasets/core/utils/py_utils.py

@@ -25,10 +25,8 @@
 import itertools
 import logging
 import os
-import random
 import re
 import shutil
-import string
 import sys
 import textwrap
 import threading
@@ -312,30 +310,6 @@ def nullcontext(enter_result: T = None) -> Iterator[T]:
   yield enter_result
 
 
-def _get_incomplete_dir(dir_name: str) -> str:
-  """Returns a temporary dir name based on `dir_name`."""
-  random_suffix = ''.join(
-      random.choice(string.ascii_uppercase + string.digits) for _ in range(6)
-  )
-  dir_name = epath.Path(dir_name)
-  return f'{dir_name.parent}/{constants.INCOMPLETE_PREFIX}{random_suffix}_{dir_name.name}/'
-
-
-@contextlib.contextmanager
-def incomplete_dir(dirname: epath.PathLike) -> 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.mkdir(parents=True, exist_ok=True)
-  try:
-    yield tmp_dir
-    tmp_path.rename(dirname)
-  finally:
-    if tmp_path.exists():
-      tmp_path.rmtree()
-
-
 def _tmp_file_name(path: epath.PathLike) -> epath.Path:
   path = epath.Path(path)
   return (