From b8e770747f7e07a2ab6f9ffe0498c52d3739833b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Saugat=20Pachhai=20=28=E0=A4=B8=E0=A5=8C=E0=A4=97=E0=A4=BE?= =?UTF-8?q?=E0=A4=A4=29?= Date: Mon, 16 May 2022 11:23:44 +0545 Subject: [PATCH] fs.utils: copy dvc.utils.fs to dvc.fs.utils --- dvc/data/db/index.py | 2 +- dvc/data/db/local.py | 5 +- dvc/data/stage.py | 2 +- dvc/data/tree.py | 2 +- dvc/fs/dvc.py | 2 +- dvc/fs/gdrive.py | 3 +- dvc/fs/generic.py | 5 +- dvc/fs/local.py | 7 +- dvc/fs/s3.py | 4 +- dvc/fs/ssh.py | 2 +- dvc/fs/utils.py | 218 ++++++++++++++++++++++++++++++++++++++++++- 11 files changed, 233 insertions(+), 19 deletions(-) diff --git a/dvc/data/db/index.py b/dvc/data/db/index.py index 21bb3b12f..77c9d7c12 100644 --- a/dvc/data/db/index.py +++ b/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) diff --git a/dvc/data/db/local.py b/dvc/data/db/local.py index ceae2439e..fbc3623a3 100644 --- a/dvc/data/db/local.py +++ b/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) diff --git a/dvc/data/stage.py b/dvc/data/stage.py index 9879e2f4a..9d15440b6 100644 --- a/dvc/data/stage.py +++ b/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 diff --git a/dvc/data/tree.py b/dvc/data/tree.py index 90466ee13..af0f4cdb5 100644 --- a/dvc/data/tree.py +++ b/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("")) diff --git a/dvc/fs/dvc.py b/dvc/fs/dvc.py index b3251f5f3..da4d5520b 100644 --- a/dvc/fs/dvc.py +++ b/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) diff --git a/dvc/fs/gdrive.py b/dvc/fs/gdrive.py index 8fdae8f0a..8f2cec267 100644 --- a/dvc/fs/gdrive.py +++ b/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" diff --git a/dvc/fs/generic.py b/dvc/fs/generic.py index 080c599b2..3c4b8fbac 100644 --- a/dvc/fs/generic.py +++ b/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( diff --git a/dvc/fs/local.py b/dvc/fs/local.py index fea355605..5160091c2 100644 --- a/dvc/fs/local.py +++ b/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) diff --git a/dvc/fs/s3.py b/dvc/fs/s3.py index 6a12d24a7..3a9c2edfa 100644 --- a/dvc/fs/s3.py +++ b/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 diff --git a/dvc/fs/ssh.py b/dvc/fs/ssh.py index 97ac12271..6d983a83a 100644 --- a/dvc/fs/ssh.py +++ b/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 diff --git a/dvc/fs/utils.py b/dvc/fs/utils.py index 5ab5aa34d..0378ae031 100644 --- a/dvc/fs/utils.py +++ b/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