fs.utils: copy dvc.utils.fs to dvc.fs.utils

dvc/data/db/index.py

@@ -89,8 +89,8 @@ class ObjectDBIndex(ObjectDBIndexBase):
         name: str,
     ):  # pylint: disable=super-init-not-called
         from dvc.fs.local import LocalFileSystem
+        from dvc.fs.utils import makedirs
         from dvc.objects.cache import Cache, Index
-        from dvc.utils.fs import makedirs
 
         self.index_dir = os.path.join(tmp_dir, self.INDEX_DIR, name)
         makedirs(self.index_dir, exist_ok=True)

dvc/data/db/local.py

@@ -7,11 +7,10 @@ from funcy import cached_property
 from shortuuid import uuid
 
 from dvc.fs.system import umask
+from dvc.fs.utils import copyfile, relpath, remove, walk_files
 from dvc.objects.db import ObjectDB, noop, wrap_iter
 from dvc.objects.errors import ObjectFormatError
 from dvc.objects.hash_info import HashInfo
-from dvc.utils import relpath
-from dvc.utils.fs import copyfile, remove, walk_files
 
 logger = logging.getLogger(__name__)
 
@@ -50,7 +49,7 @@ class LocalObjectDB(ObjectDB):
         os.chmod(to_info, self._file_mode)
 
     def makedirs(self, fs_path):
-        from dvc.utils.fs import makedirs
+        from dvc.fs.utils import makedirs
 
         makedirs(fs_path, exist_ok=True, mode=self._dir_mode)
 

dvc/data/stage.py

@@ -31,7 +31,7 @@ _STAGING_MEMFS_PATH = "dvc-staging"
 
 def _upload_file(from_fs_path, fs, odb, upload_odb, callback=None):
     from dvc.fs._callback import FsspecCallback
-    from dvc.utils import tmp_fname
+    from dvc.fs.utils import tmp_fname
 
     from .stream import HashedStreamReader
 

dvc/data/tree.py

@@ -56,7 +56,7 @@ class Tree(HashFile):
 
     def digest(self, hash_info: Optional["HashInfo"] = None):
         from dvc.fs.memory import MemoryFileSystem
-        from dvc.utils import tmp_fname
+        from dvc.fs.utils import tmp_fname
 
         memfs = MemoryFileSystem()
         fs_path = "memory://{}".format(tmp_fname(""))

dvc/fs/dvc.py

@@ -468,7 +468,7 @@ class DvcFileSystem(FileSystem):
         return self.fs.config
 
     def from_os_path(self, path):
-        from dvc.utils import as_posix
+        from .utils import as_posix
 
         if os.path.isabs(path):
             path = os.path.relpath(path, self.repo.root_dir)

dvc/fs/gdrive.py

@@ -6,9 +6,10 @@ from funcy import cached_property, wrap_prop
 
 from dvc.exceptions import DvcException
 from dvc.scheme import Schemes
-from dvc.utils import format_link, tmp_fname
+from dvc.utils import format_link
 
 from .base import FileSystem
+from .utils import tmp_fname
 
 logger = logging.getLogger(__name__)
 FOLDER_MIME_TYPE = "application/vnd.google-apps.folder"

dvc/fs/generic.py

@@ -3,11 +3,10 @@ import logging
 import os
 from typing import TYPE_CHECKING, List, Optional
 
-from dvc.utils.fs import as_atomic
-
 from ._callback import DEFAULT_CALLBACK
 from .base import RemoteActionNotImplemented
 from .local import LocalFileSystem
+from .utils import as_atomic
 
 if TYPE_CHECKING:
     from ._callback import FsspecCallback
@@ -154,7 +153,7 @@ def test_links(
     to_fs: "FileSystem",
     to_path: "AnyFSPath",
 ) -> List["AnyFSPath"]:
-    from dvc.utils import tmp_fname
+    from .utils import tmp_fname
 
     from_file = from_fs.path.join(from_path, tmp_fname())
     to_file = to_fs.path.join(

dvc/fs/local.py

@@ -1,5 +1,6 @@
 import logging
 import os
+import shutil
 import threading
 
 from fsspec import AbstractFileSystem
@@ -7,10 +8,9 @@ from funcy import cached_property, wrap_prop
 
 from dvc.fs import system
 from dvc.scheme import Schemes
-from dvc.utils import tmp_fname
-from dvc.utils.fs import copy_fobj_to_file, copyfile, makedirs, move, remove
 
 from .base import FileSystem
+from .utils import copyfile, makedirs, move, remove, tmp_fname
 
 logger = logging.getLogger(__name__)
 
@@ -188,7 +188,8 @@ class LocalFileSystem(FileSystem):
         self.makedirs(self.path.parent(to_info))
         tmp_info = self.path.join(self.path.parent(to_info), tmp_fname(""))
         try:
-            copy_fobj_to_file(fobj, tmp_info)
+            with open(tmp_info, "wb+") as fdest:
+                shutil.copyfileobj(fobj, fdest)
             os.rename(tmp_info, to_info)
         except Exception:
             self.remove(tmp_info)

dvc/fs/s3.py

@@ -38,14 +38,14 @@ class S3FileSystem(ObjectFileSystem):
 
         from boto3.s3.transfer import TransferConfig
 
-        from dvc.utils import conversions
+        from .utils import human_readable_to_bytes
 
         config, transfer_config = {}, {}
         for key, value in s3_config.items():
             if key in self._TRANSFER_CONFIG_ALIASES:
                 if key in {"multipart_chunksize", "multipart_threshold"}:
                     # cast human readable sizes (like 24MiB) to integers
-                    value = conversions.human_readable_to_bytes(value)
+                    value = human_readable_to_bytes(value)
                 else:
                     value = int(value)
                 transfer_config[self._TRANSFER_CONFIG_ALIASES[key]] = value

dvc/fs/ssh.py

@@ -6,10 +6,10 @@ from funcy import cached_property, memoize, silent, wrap_prop, wrap_with
 
 from dvc import prompt
 from dvc.scheme import Schemes
-from dvc.utils.fs import as_atomic
 
 from ._callback import DEFAULT_CALLBACK
 from .base import FileSystem
+from .utils import as_atomic
 
 DEFAULT_PORT = 22
 

dvc/fs/utils.py

@@ -1,9 +1,23 @@
+import errno
+import logging
 import os
+import shutil
 import stat
-from typing import TYPE_CHECKING
+import sys
+from contextlib import contextmanager, suppress
+from typing import TYPE_CHECKING, Iterator
+
+from . import system
 
 if TYPE_CHECKING:
-    from .base import AnyFSPath
+    from ._callback import FsspecCallback
+    from .base import AnyFSPath, FileSystem
+
+
+logger = logging.getLogger(__name__)
+
+
+LOCAL_CHUNK_SIZE = 2**20  # 1 MB
 
 
 def is_exec(mode: int) -> bool:
@@ -29,3 +43,203 @@ def relpath(path: "AnyFSPath", start: "AnyFSPath" = os.curdir) -> "AnyFSPath":
         if not os.path.commonprefix([start, path]):
             return path
     return os.path.relpath(path, start)
+
+
+def move(src: "AnyFSPath", dst: "AnyFSPath") -> None:
+    """Atomically move src to dst and chmod it with mode.
+
+    Moving is performed in two stages to make the whole operation atomic in
+    case src and dst are on different filesystems and actual physical copying
+    of data is happening.
+    """
+    from shortuuid import uuid
+
+    dst = os.path.abspath(dst)
+    tmp = f"{dst}.{uuid()}"
+
+    if os.path.islink(src):
+        shutil.copy(src, tmp)
+        _unlink(src, _chmod)
+    else:
+        shutil.move(src, tmp)
+
+    shutil.move(tmp, dst)
+
+
+def _chmod(func, p, excinfo):  # pylint: disable=unused-argument
+    perm = os.lstat(p).st_mode
+    perm |= stat.S_IWRITE
+
+    try:
+        os.chmod(p, perm)
+    except OSError as exc:
+        # broken symlink or file is not owned by us
+        if exc.errno not in [errno.ENOENT, errno.EPERM]:
+            raise
+
+    func(p)
+
+
+def _unlink(path: "AnyFSPath", onerror):
+    try:
+        os.unlink(path)
+    except OSError:
+        onerror(os.unlink, path, sys.exc_info())
+
+
+def remove(path: "AnyFSPath") -> None:
+    logger.debug("Removing '%s'", path)
+
+    try:
+        if os.path.isdir(path):
+            shutil.rmtree(path, onerror=_chmod)
+        else:
+            _unlink(path, _chmod)
+    except OSError as exc:
+        if exc.errno != errno.ENOENT:
+            raise
+
+
+def makedirs(path, exist_ok: bool = False, mode: int = None) -> None:
+    if mode is None:
+        os.makedirs(path, exist_ok=exist_ok)
+        return
+
+    # Modified version of os.makedirs() with support for extended mode
+    # (e.g. S_ISGID)
+    head, tail = os.path.split(path)
+    if not tail:
+        head, tail = os.path.split(head)
+    if head and tail and not os.path.exists(head):
+        try:
+            makedirs(head, exist_ok=exist_ok, mode=mode)
+        except FileExistsError:
+            # Defeats race condition when another thread created the path
+            pass
+        cdir = os.curdir
+        if isinstance(tail, bytes):
+            cdir = bytes(os.curdir, "ASCII")  # type: ignore[assignment]
+        if tail == cdir:  # xxx/newdir/. exists if xxx/newdir exists
+            return
+    try:
+        os.mkdir(path, mode)
+    except OSError:
+        # Cannot rely on checking for EEXIST, since the operating system
+        # could give priority to other errors like EACCES or EROFS
+        if not exist_ok or not os.path.isdir(path):
+            raise
+
+    try:
+        os.chmod(path, mode)
+    except OSError:
+        logger.trace(  # type: ignore[attr-defined]
+            "failed to chmod '%o' '%s'", mode, path, exc_info=True
+        )
+
+
+def copyfile(
+    src: "AnyFSPath",
+    dest: "AnyFSPath",
+    callback: "FsspecCallback" = None,
+    no_progress_bar: bool = False,
+    name: str = None,
+) -> None:
+    """Copy file with progress bar"""
+    name = name if name else os.path.basename(dest)
+    total = os.stat(src).st_size
+
+    if os.path.isdir(dest):
+        dest = os.path.join(dest, os.path.basename(src))
+
+    if callback:
+        callback.set_size(total)
+
+    try:
+        system.reflink(src, dest)
+    except OSError:
+        from dvc.fs._callback import FsspecCallback
+
+        with open(src, "rb") as fsrc, open(dest, "wb+") as fdest:
+            with FsspecCallback.as_tqdm_callback(
+                callback,
+                size=total,
+                bytes=True,
+                disable=no_progress_bar,
+                desc=name,
+            ) as cb:
+                wrapped = cb.wrap_attr(fdest, "write")
+                while True:
+                    buf = fsrc.read(LOCAL_CHUNK_SIZE)
+                    if not buf:
+                        break
+                    wrapped.write(buf)
+
+    if callback:
+        callback.absolute_update(total)
+
+
+def walk_files(directory: "AnyFSPath") -> Iterator["AnyFSPath"]:
+    for root, _, files in os.walk(directory):
+        for f in files:
+            yield os.path.join(root, f)
+
+
+def tmp_fname(fname: "AnyFSPath" = "") -> "AnyFSPath":
+    """Temporary name for a partial download"""
+    from shortuuid import uuid
+
+    return os.fspath(fname) + "." + uuid() + ".tmp"
+
+
+def as_posix(path: str) -> str:
+    import ntpath
+    import posixpath
+
+    return path.replace(ntpath.sep, posixpath.sep)
+
+
+@contextmanager
+def as_atomic(
+    fs: "FileSystem", to_info: "AnyFSPath", create_parents: bool = False
+) -> Iterator["AnyFSPath"]:
+    parent = fs.path.parent(to_info)
+    if create_parents:
+        fs.makedirs(parent, exist_ok=True)
+
+    tmp_info = fs.path.join(parent, tmp_fname())
+    try:
+        yield tmp_info
+    except BaseException:
+        # Handle stuff like KeyboardInterrupt
+        # as well as other errors that might
+        # arise during file transfer.
+        with suppress(FileNotFoundError):
+            fs.remove(tmp_info)
+        raise
+    else:
+        fs.move(tmp_info, to_info)
+
+
+# https://github.com/aws/aws-cli/blob/5aa599949f60b6af554fd5714d7161aa272716f7/awscli/customizations/s3/utils.py
+MULTIPLIERS = {
+    "kb": 1024,
+    "mb": 1024**2,
+    "gb": 1024**3,
+    "tb": 1024**4,
+    "kib": 1024,
+    "mib": 1024**2,
+    "gib": 1024**3,
+    "tib": 1024**4,
+}
+
+
+def human_readable_to_bytes(value: str) -> int:
+    value = value.lower()
+    suffix = ""
+    if value.endswith(tuple(MULTIPLIERS.keys())):
+        size = 2
+        size += value[-2] == "i"  # KiB, MiB etc
+        value, suffix = value[:-size], value[-size:]
+
+    multiplier = MULTIPLIERS.get(suffix, 1)
+    return int(value) * multiplier