Add option to write incomplete files to a subfolder

This can help with getting the dataset folder clean

PiperOrigin-RevId: 695237951

Author: tomvdw

tensorflow_datasets/core/utils/py_utils.py

@@ -308,17 +308,32 @@ def _tmp_file_prefix() -> str:
   return f'{constants.INCOMPLETE_PREFIX}{uuid.uuid4().hex}'
 
 
-def _tmp_file_name(path: epath.PathLike) -> epath.Path:
+def _tmp_file_name(
+    path: epath.PathLike,
+    subfolder: str | None = None,
+) -> epath.Path:
+  """Returns the temporary file name for the given path.
+
+  Args:
+    path: The path to the file.
+    subfolder: The subfolder to use. If None, then the parent of the path will
+      be used.
+  """
   path = epath.Path(path)
-  return path.parent / f'{_tmp_file_prefix()}.{path.name}'
+  file_name = f'{_tmp_file_prefix()}.{path.name}'
+  if subfolder:
+    return path.parent / subfolder / file_name
+  else:
+    return path.parent / file_name
 
 
 @contextlib.contextmanager
 def incomplete_file(
     path: epath.Path,
+    subfolder: str | None = None,
 ) -> Iterator[epath.Path]:
   """Writes to path atomically, by writing to temp file and renaming it."""
-  tmp_path = _tmp_file_name(path)
+  tmp_path = _tmp_file_name(path, subfolder=subfolder)
   try:
     yield tmp_path
     tmp_path.replace(path)

tensorflow_datasets/core/utils/py_utils_test.py

@@ -14,7 +14,9 @@
 # limitations under the License.
 
 import collections
+import os
 import pathlib
+from unittest import mock
 
 from etils import epath
 import pytest
@@ -370,5 +372,17 @@ def test_make_valid_name(name: str, expected: str):
   assert py_utils.make_valid_name(name) == expected
 
 
+@pytest.mark.parametrize(
+    ['path', 'subfolder', 'expected'],
+    [
+        ('/a/file.ext', None, '/a/foobar.file.ext'),
+        ('/a/file.ext', 'sub', '/a/sub/foobar.file.ext'),
+    ],
+)
+def test_tmp_file_name(path, subfolder, expected):
+  with mock.patch.object(py_utils, '_tmp_file_prefix', return_value='foobar'):
+    assert os.fspath(py_utils._tmp_file_name(path, subfolder)) == expected
+
+
 if __name__ == '__main__':
   tf.test.main()