diff --git a/dvc/checkout.py b/dvc/checkout.py index c69f58e64..6a827e64a 100644 --- a/dvc/checkout.py +++ b/dvc/checkout.py @@ -35,7 +35,9 @@ def _changed(path_info, fs, obj, cache, state=None): actual = state.get(path_info, fs) if state else None if actual is None: try: - actual = stage(cache, path_info, fs, obj.hash_info.name).hash_info + actual = stage( + cache, path_info, fs, obj.hash_info.name, dry_run=True + )[1].hash_info except FileNotFoundError: logger.debug("'%s' doesn't exist.", path_info) return True @@ -61,7 +63,9 @@ def _remove(path_info, fs, cache, force=False): fs.remove(path_info) return - current = stage(cache, path_info, fs, fs.PARAM_CHECKSUM).hash_info + current = stage(cache, path_info, fs, fs.PARAM_CHECKSUM, dry_run=True)[ + 1 + ].hash_info try: obj = load(cache, current) check(cache, obj) diff --git a/dvc/data_cloud.py b/dvc/data_cloud.py index 76170b002..c9c394f5c 100644 --- a/dvc/data_cloud.py +++ b/dvc/data_cloud.py @@ -6,8 +6,8 @@ from typing import TYPE_CHECKING, Iterable, Optional from dvc.objects.db import get_index if TYPE_CHECKING: + from dvc.hash_info import HashInfo from dvc.objects.db.base import ObjectDB - from dvc.objects.file import HashFile logger = logging.getLogger(__name__) @@ -64,7 +64,7 @@ class DataCloud: def push( self, - objs: Iterable["HashFile"], + objs: Iterable["HashInfo"], jobs: Optional[int] = None, remote: Optional[str] = None, odb: Optional["ObjectDB"] = None, @@ -93,7 +93,7 @@ class DataCloud: def pull( self, - objs: Iterable["HashFile"], + objs: Iterable["HashInfo"], jobs: Optional[int] = None, remote: Optional[str] = None, odb: Optional["ObjectDB"] = None, @@ -123,7 +123,7 @@ class DataCloud: def status( self, - objs: Iterable["HashFile"], + objs: Iterable["HashInfo"], jobs: Optional[int] = None, remote: Optional[str] = None, odb: Optional["ObjectDB"] = None, diff --git a/dvc/dependency/repo.py b/dvc/dependency/repo.py index 234749900..c55217603 100644 --- a/dvc/dependency/repo.py +++ b/dvc/dependency/repo.py @@ -1,8 +1,8 @@ import os from collections import defaultdict +from copy import copy from typing import TYPE_CHECKING, Dict, Optional, Set, Tuple -from funcy import first from voluptuous import Required from dvc.path_info import PathInfo @@ -10,6 +10,7 @@ from dvc.path_info import PathInfo from .base import Dependency if TYPE_CHECKING: + from dvc.hash_info import HashInfo from dvc.objects.db.base import ObjectDB from dvc.objects.file import HashFile @@ -91,23 +92,24 @@ class RepoDependency(Dependency): def get_used_objs( self, **kwargs - ) -> Dict[Optional["ObjectDB"], Set["HashFile"]]: + ) -> Dict[Optional["ObjectDB"], Set["HashInfo"]]: used, _ = self._get_used_and_obj(**kwargs) return used def _get_used_and_obj( self, obj_only=False, **kwargs - ) -> Tuple[Dict[Optional["ObjectDB"], Set["HashFile"]], "HashFile"]: + ) -> Tuple[Dict[Optional["ObjectDB"], Set["HashInfo"]], "HashFile"]: from dvc.config import NoRemoteError from dvc.exceptions import NoOutputOrStageError, PathMissingError - from dvc.objects.stage import get_staging, stage + from dvc.objects.stage import stage + from dvc.objects.tree import Tree local_odb = self.repo.odb.local locked = kwargs.pop("locked", True) with self._make_repo( locked=locked, cache_dir=local_odb.cache_dir ) as repo: - used_objs = defaultdict(set) + used_obj_ids = defaultdict(set) rev = repo.get_rev() if locked and self.def_repo.get(self.PARAM_REV_LOCK) is None: self.def_repo[self.PARAM_REV_LOCK] = rev @@ -115,7 +117,7 @@ class RepoDependency(Dependency): path_info = PathInfo(repo.root_dir) / str(self.def_path) if not obj_only: try: - for odb, objs in repo.used_objs( + for odb, obj_ids in repo.used_objs( [os.fspath(path_info)], force=True, jobs=kwargs.get("jobs"), @@ -124,14 +126,14 @@ class RepoDependency(Dependency): if odb is None: odb = repo.cloud.get_remote_odb() odb.read_only = True - self._check_circular_import(objs) - used_objs[odb].update(objs) + self._check_circular_import(odb, obj_ids) + used_obj_ids[odb].update(obj_ids) except (NoRemoteError, NoOutputOrStageError): pass try: - staged_obj = stage( - None, + staging, staged_obj = stage( + local_odb, path_info, repo.repo_fs, local_odb.fs.PARAM_CHECKSUM, @@ -140,27 +142,51 @@ class RepoDependency(Dependency): raise PathMissingError( self.def_path, self.def_repo[self.PARAM_URL] ) from exc - staging = get_staging() + staging = copy(staging) staging.read_only = True self._staged_objs[rev] = staged_obj - used_objs[staging].add(staged_obj) - return used_objs, staged_obj + used_obj_ids[staging].add(staged_obj.hash_info) + if isinstance(staged_obj, Tree): + used_obj_ids[staging].update( + entry.hash_info for _, entry in staged_obj + ) + return used_obj_ids, staged_obj - def _check_circular_import(self, objs): + def _check_circular_import(self, odb, obj_ids): from dvc.exceptions import CircularImportError from dvc.fs.repo import RepoFileSystem + from dvc.objects.db.reference import ReferenceObjectDB from dvc.objects.tree import Tree - obj = first(objs) - if isinstance(obj, Tree): - _, obj = first(obj) - if not isinstance(obj.fs, RepoFileSystem): + if not isinstance(odb, ReferenceObjectDB): return - self_url = self.repo.url or self.repo.root_dir - if obj.fs.repo_url is not None and obj.fs.repo_url == self_url: - raise CircularImportError(self, obj.fs.repo_url, self_url) + def iter_objs(): + for hash_info in obj_ids: + if hash_info.isdir: + tree = Tree.load(odb, hash_info) + yield from (odb.get(entry.hash_info) for _, entry in tree) + else: + yield odb.get(hash_info) + + checked_urls = set() + for obj in iter_objs(): + if not isinstance(obj.fs, RepoFileSystem): + continue + if ( + obj.fs.repo_url in checked_urls + or obj.fs.root_dir in checked_urls + ): + continue + self_url = self.repo.url or self.repo.root_dir + if ( + obj.fs.repo_url is not None + and obj.fs.repo_url == self_url + or obj.fs.root_dir == self.repo.root_dir + ): + raise CircularImportError(self, obj.fs.repo_url, self_url) + checked_urls.update([obj.fs.repo_url, obj.fs.root_dir]) def get_obj(self, filter_info=None, **kwargs): locked = kwargs.get("locked", True) diff --git a/dvc/external_repo.py b/dvc/external_repo.py index b693e7953..88ec1e1ec 100644 --- a/dvc/external_repo.py +++ b/dvc/external_repo.py @@ -42,14 +42,6 @@ def external_repo( config = _get_remote_config(url) if os.path.isdir(url) else {} config.update(cache_config) - def make_repo(path, **_kwargs): - _config = cache_config.copy() - if os.path.isdir(url): - rel = os.path.relpath(path, _kwargs["scm"].root_dir) - repo_path = os.path.join(url, rel) - _config.update(_get_remote_config(repo_path)) - return Repo(path, config=_config, **_kwargs) - root_dir = path if for_write else os.path.realpath(path) repo_kwargs = dict( root_dir=root_dir, @@ -57,7 +49,7 @@ def external_repo( scm=None if for_write else Git(root_dir), rev=None if for_write else rev, config=config, - repo_factory=make_repo, + repo_factory=erepo_factory(url, cache_config), **kwargs, ) @@ -87,6 +79,18 @@ def external_repo( _remove(path) +def erepo_factory(url, cache_config, *args, **kwargs): + def make_repo(path, **_kwargs): + _config = cache_config.copy() + if os.path.isdir(url): + rel = os.path.relpath(path, _kwargs["scm"].root_dir) + repo_path = os.path.join(url, rel) + _config.update(_get_remote_config(repo_path)) + return Repo(path, config=_config, **_kwargs) + + return make_repo + + CLONES: Dict[str, str] = {} CACHE_DIRS: Dict[str, str] = {} diff --git a/dvc/fs/__init__.py b/dvc/fs/__init__.py index 4aee14a9b..230a78733 100644 --- a/dvc/fs/__init__.py +++ b/dvc/fs/__init__.py @@ -32,8 +32,9 @@ FS_MAP = { } -def get_fs_cls(remote_conf): - scheme = urlparse(remote_conf["url"]).scheme +def get_fs_cls(remote_conf, scheme=None): + if not scheme: + scheme = urlparse(remote_conf["url"]).scheme return FS_MAP.get(scheme, LocalFileSystem) diff --git a/dvc/fs/base.py b/dvc/fs/base.py index 99a06e929..33a6d1fa7 100644 --- a/dvc/fs/base.py +++ b/dvc/fs/base.py @@ -59,6 +59,11 @@ class BaseFileSystem: self.jobs = kwargs.get("jobs") or self._JOBS self.hash_jobs = kwargs.get("checksum_jobs") or self.HASH_JOBS + self._config = kwargs + + @property + def config(self): + return self._config @classmethod def _strip_protocol(cls, path: str): diff --git a/dvc/fs/dvc.py b/dvc/fs/dvc.py index 612b40163..2e5182e4e 100644 --- a/dvc/fs/dvc.py +++ b/dvc/fs/dvc.py @@ -32,6 +32,10 @@ class DvcFileSystem(BaseFileSystem): # pylint:disable=abstract-method super().__init__(**kwargs) self.repo = kwargs["repo"] + @property + def config(self): + raise NotImplementedError + def _find_outs(self, path, *args, **kwargs): outs = self.repo.find_outs_by_path(path, *args, **kwargs) diff --git a/dvc/fs/repo.py b/dvc/fs/repo.py index 37e75ec09..10f576451 100644 --- a/dvc/fs/repo.py +++ b/dvc/fs/repo.py @@ -35,14 +35,29 @@ class RepoFileSystem(BaseFileSystem): # pylint:disable=abstract-method scheme = "local" PARAM_CHECKSUM = "md5" + PARAM_REPO_URL = "repo_url" + PARAM_REPO_ROOT = "repo_root" + PARAM_REV = "rev" + PARAM_CACHE_DIR = "cache_dir" + PARAM_CACHE_TYPES = "cache_types" + PARAM_SUBREPOS = "subrepos" def __init__( - self, repo=None, subrepos=False, repo_factory: RepoFactory = None + self, + repo: Optional["Repo"] = None, + subrepos=False, + repo_factory: RepoFactory = None, + **kwargs, ): super().__init__() from dvc.utils.collections import PathStringTrie + if repo is None: + repo, repo_factory = self._repo_from_fs_config( + subrepos=subrepos, **kwargs + ) + if not repo_factory: from dvc.repo import Repo @@ -72,6 +87,67 @@ class RepoFileSystem(BaseFileSystem): # pylint:disable=abstract-method return None return self._main_repo.url + @property + def config(self): + return { + self.PARAM_REPO_URL: self.repo_url, + self.PARAM_REPO_ROOT: self.root_dir, + self.PARAM_REV: getattr(self._main_repo.fs, "rev", None), + self.PARAM_CACHE_DIR: os.path.abspath( + self._main_repo.odb.local.cache_dir + ), + self.PARAM_CACHE_TYPES: self._main_repo.odb.local.cache_types, + self.PARAM_SUBREPOS: self._traverse_subrepos, + } + + @classmethod + def _repo_from_fs_config( + cls, **config + ) -> Tuple["Repo", Optional["RepoFactory"]]: + from dvc.external_repo import erepo_factory, external_repo + from dvc.repo import Repo + + url = config.get(cls.PARAM_REPO_URL) + root = config.get(cls.PARAM_REPO_ROOT) + assert url or root + + def _open(*args, **kwargs): + # NOTE: if original repo was an erepo (and has a URL), + # we cannot use Repo.open() since it will skip erepo + # cache/remote setup for local URLs + if url is None: + return Repo.open(*args, **kwargs) + return external_repo(*args, **kwargs) + + cache_dir = config.get(cls.PARAM_CACHE_DIR) + cache_config = ( + {} + if not cache_dir + else { + "cache": { + "dir": cache_dir, + "type": config.get(cls.PARAM_CACHE_TYPES), + } + } + ) + repo_kwargs: dict = { + "rev": config.get(cls.PARAM_REV), + "subrepos": config.get(cls.PARAM_SUBREPOS, False), + "uninitialized": True, + } + factory: Optional["RepoFactory"] = None + if url is None: + repo_kwargs["config"] = cache_config + else: + repo_kwargs["cache_dir"] = cache_dir + factory = erepo_factory(url, cache_config) + + with _open( + url if url else root, + **repo_kwargs, + ) as repo: + return repo, factory + def _get_repo(self, path: str) -> Optional["Repo"]: """Returns repo that the path falls in, using prefix. diff --git a/dvc/hash_info.py b/dvc/hash_info.py index 72fdbe6f0..d12dd52a4 100644 --- a/dvc/hash_info.py +++ b/dvc/hash_info.py @@ -14,6 +14,7 @@ class HashInfo: value: Optional[str] size: Optional[int] = field(default=None, compare=False) nfiles: Optional[int] = field(default=None, compare=False) + obj_name: Optional[str] = field(default=None, compare=False) def __bool__(self): return bool(self.value) diff --git a/dvc/objects/__init__.py b/dvc/objects/__init__.py index bfe1321d4..ad9721a46 100644 --- a/dvc/objects/__init__.py +++ b/dvc/objects/__init__.py @@ -1,8 +1,5 @@ import logging -from concurrent.futures import ThreadPoolExecutor, as_completed -from typing import TYPE_CHECKING, Optional - -from dvc.progress import Tqdm +from typing import TYPE_CHECKING from .tree import Tree @@ -15,50 +12,12 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) -def save( - odb: "ObjectDB", - obj: "HashFile", - jobs: Optional[int] = None, - **kwargs, -): - if isinstance(obj, Tree): - from .stage import get_staging - - with ThreadPoolExecutor(max_workers=jobs) as executor: - for future in Tqdm( - as_completed( - executor.submit( - odb.add, - entry.path_info, - entry.fs, - entry.hash_info, - **kwargs, - ) - for _, entry in obj - ), - desc="Saving files", - total=len(obj), - unit="file", - ): - future.result() - - # if dir cache for this tree has already been staged, move the - # staged object when possible - staging = get_staging(odb) - if staging.exists(obj.hash_info): - obj = staging.get(obj.hash_info) - kwargs["move"] = True - - assert obj.path_info and obj.fs and obj.hash_info - odb.add(obj.path_info, obj.fs, obj.hash_info, **kwargs) - - -def check(odb: "ObjectDB", obj: "HashFile"): +def check(odb: "ObjectDB", obj: "HashFile", **kwargs): if isinstance(obj, Tree): for _, entry in obj: - odb.check(entry.hash_info) + odb.check(entry.hash_info, **kwargs) - odb.check(obj.hash_info) + odb.check(obj.hash_info, **kwargs) def load(odb: "ObjectDB", hash_info: "HashInfo") -> "HashFile": diff --git a/dvc/objects/db/base.py b/dvc/objects/db/base.py index 0a2a0592b..3f1e3face 100644 --- a/dvc/objects/db/base.py +++ b/dvc/objects/db/base.py @@ -13,7 +13,7 @@ from dvc.progress import Tqdm if TYPE_CHECKING: from dvc.fs.base import BaseFileSystem from dvc.hash_info import HashInfo - from dvc.types import AnyPath + from dvc.types import AnyPath, DvcPath logger = logging.getLogger(__name__) @@ -76,7 +76,14 @@ class ObjectDB: hash_info, ) - def _add_file(self, from_fs, from_info, to_info, move): + def _add_file( + self, + from_fs: "BaseFileSystem", + from_info: "AnyPath", + to_info: "DvcPath", + _hash_info: "HashInfo", + move: bool = False, + ): self.makedirs(to_info.parent) use_move = isinstance(from_fs, type(self.fs)) and move try: @@ -131,8 +138,8 @@ class ObjectDB: pass cache_info = self.hash_to_path_info(hash_info.value) - # using our makedirs to create dirs with proper permissions - self._add_file(fs, path_info, cache_info, move) + self._add_file(fs, path_info, cache_info, hash_info, move) + try: if verify: self.check(hash_info, check_hash=True) @@ -145,7 +152,7 @@ class ObjectDB: if callback: callback(1) - def hash_to_path_info(self, hash_): + def hash_to_path_info(self, hash_) -> "DvcPath": return self.path_info / hash_[0:2] / hash_[2:] # Override to return path as a string instead of PathInfo for clouds @@ -199,9 +206,10 @@ class ObjectDB: self.fs.remove(obj.path_info) raise - # making cache file read-only so we don't need to check it - # next time - self.protect(obj.path_info) + if check_hash: + # making cache file read-only so we don't need to check it + # next time + self.protect(obj.path_info) def _list_paths(self, prefix=None, progress_callback=None): if prefix: @@ -378,17 +386,20 @@ class ObjectDB: def _remove_unpacked_dir(self, hash_): pass - def gc(self, used, jobs=None): + def gc(self, used, jobs=None, cache_odb=None, shallow=True): from ..tree import Tree if self.read_only: raise ObjectDBPermissionError("Cannot gc read-only ODB") + if not cache_odb: + cache_odb = self used_hashes = set() - for obj in used: - used_hashes.add(obj.hash_info.value) - if isinstance(obj, Tree): + for hash_info in used: + used_hashes.add(hash_info.value) + if hash_info.isdir and not shallow: + tree = Tree.load(cache_odb, hash_info) used_hashes.update( - entry_obj.hash_info.value for _, entry_obj in obj + entry_obj.hash_info.value for _, entry_obj in tree ) removed = False diff --git a/dvc/objects/db/reference.py b/dvc/objects/db/reference.py new file mode 100644 index 000000000..12e46b8ba --- /dev/null +++ b/dvc/objects/db/reference.py @@ -0,0 +1,77 @@ +import io +import logging +import os +from typing import TYPE_CHECKING, Dict + +from dvc.scheme import Schemes + +from ..errors import ObjectFormatError +from ..reference import ReferenceHashFile +from .base import ObjectDB + +if TYPE_CHECKING: + from dvc.fs.base import BaseFileSystem + from dvc.hash_info import HashInfo + from dvc.types import AnyPath, DvcPath + +logger = logging.getLogger(__name__) + + +class ReferenceObjectDB(ObjectDB): + """Reference ODB. + + File objects are stored as ReferenceHashFiles which reference paths outside + of the staging ODB fs. Tree objects are stored natively. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._fs_cache: Dict[tuple, "BaseFileSystem"] = {} + + def get(self, hash_info: "HashInfo"): + if hash_info.isdir: + return super().get(hash_info) + path_info = self.hash_to_path(hash_info.value) + try: + with self.fs.open(path_info, "rb") as fobj: + ref_file = ReferenceHashFile.from_bytes( + fobj.read(), fs_cache=self._fs_cache + ) + except OSError: + raise FileNotFoundError + try: + ref_file.check(self, check_hash=False) + except ObjectFormatError: + self.fs.remove(path_info) + raise + return ref_file + + def _add_file( + self, + from_fs: "BaseFileSystem", + from_info: "AnyPath", + to_info: "DvcPath", + hash_info: "HashInfo", + move: bool = False, + ): + self.makedirs(to_info.parent) + if hash_info.isdir: + return super()._add_file( + from_fs, from_info, to_info, hash_info, move + ) + ref_file = ReferenceHashFile(from_info, from_fs, hash_info) + ref_fobj = io.BytesIO(ref_file.to_bytes()) + ref_fobj.seek(0) + try: + self.fs.upload_fobj(ref_fobj, to_info) + except OSError as exc: + if isinstance(exc, FileExistsError) or ( + os.name == "nt" + and exc.__context__ + and isinstance(exc.__context__, FileExistsError) + ): + logger.debug("'%s' file already exists, skipping", to_info) + else: + raise + if from_fs.scheme != Schemes.LOCAL: + self._fs_cache[ReferenceHashFile.config_tuple(from_fs)] = from_fs diff --git a/dvc/objects/file.py b/dvc/objects/file.py index 4fcf6fd5f..2e4fc4300 100644 --- a/dvc/objects/file.py +++ b/dvc/objects/file.py @@ -8,7 +8,9 @@ from .errors import ObjectFormatError if TYPE_CHECKING: from dvc.fs.base import BaseFileSystem from dvc.hash_info import HashInfo - from dvc.types import DvcPath + from dvc.types import AnyPath + + from .db.base import ObjectDB logger = logging.getLogger(__name__) @@ -16,7 +18,7 @@ logger = logging.getLogger(__name__) class HashFile: def __init__( self, - path_info: Optional["DvcPath"], + path_info: Optional["AnyPath"], fs: Optional["BaseFileSystem"], hash_info: "HashInfo", name: Optional[str] = None, @@ -57,10 +59,9 @@ class HashFile: ) ) - def check(self, odb, check_hash=True): - from .stage import get_file_hash - + def check(self, odb: "ObjectDB", check_hash: bool = True): if not check_hash: + assert self.fs if not self.fs.exists(self.path_info): raise FileNotFoundError( errno.ENOENT, os.strerror(errno.ENOENT), self.path_info @@ -68,6 +69,11 @@ class HashFile: else: return None + self._check_hash(odb) + + def _check_hash(self, odb): + from .stage import get_file_hash + actual = get_file_hash( self.path_info, self.fs, self.hash_info.name, odb.state ) diff --git a/dvc/objects/reference.py b/dvc/objects/reference.py new file mode 100644 index 000000000..df6807917 --- /dev/null +++ b/dvc/objects/reference.py @@ -0,0 +1,135 @@ +import errno +import logging +import os +import pickle +from typing import TYPE_CHECKING, Optional, Tuple + +from .errors import ObjectFormatError +from .file import HashFile + +if TYPE_CHECKING: + from dvc.fs.base import BaseFileSystem + from dvc.hash_info import HashInfo + from dvc.types import AnyPath + + from .db.base import ObjectDB + +logger = logging.getLogger(__name__) + + +class ReferenceHashFile(HashFile): + PARAM_PATH = "path" + PARAM_HASH = "hash" + PARAM_MTIME = "mtime" + PARAM_SIZE = "size" + PARAM_FS_CONFIG = "fs_config" + + def __init__( + self, + path_info: "AnyPath", + fs: "BaseFileSystem", + hash_info: "HashInfo", + mtime: Optional[str] = None, + size: Optional[int] = None, + **kwargs, + ): + super().__init__(path_info, fs, hash_info, **kwargs) + cur_mtime, cur_size = self._get_mtime_and_size() + self.mtime = cur_mtime if mtime is None else mtime + self.hash_info.size = cur_size if size is None else size + + def __str__(self): + return f"ref object {self.hash_info} -> {self.path_info}" + + def check(self, odb: "ObjectDB", check_hash: bool = True): + assert self.fs + if not self.fs.exists(self.path_info): + raise FileNotFoundError( + errno.ENOENT, os.strerror(errno.ENOENT), self.path_info + ) + if not (self.mtime, self.size) == self._get_mtime_and_size(): + raise ObjectFormatError(f"{self} is changed") + if check_hash: + self._check_hash(odb) + + def _get_mtime_and_size(self) -> Tuple[Optional[str], int]: + from dvc.utils.fs import get_mtime_and_size + + assert self.fs + if hasattr(self.fs, "stat"): + try: + return get_mtime_and_size(self.path_info, self.fs) + except FileNotFoundError: + pass + return None, self.fs.getsize(self.path_info) + + def to_bytes(self): + from dvc.fs.repo import RepoFileSystem + from dvc.path_info import PathInfo + + # NOTE: dumping reference FS's this way is insecure, as the + # fully parsed remote FS config will include credentials + # + # ReferenceHashFiles should currently only be serialized in + # memory and not to disk + path_info = self.path_info + if isinstance(self.fs, RepoFileSystem): + path_info = PathInfo(path_info).relative_to(self.fs.root_dir) + + dict_ = { + self.PARAM_PATH: path_info, + self.PARAM_HASH: self.hash_info, + self.PARAM_MTIME: self.mtime, + self.PARAM_SIZE: self.size, + self.PARAM_FS_CONFIG: self.config_tuple(self.fs), + } + try: + return pickle.dumps(dict_) + except pickle.PickleError as exc: + raise ObjectFormatError(f"Could not pickle {self}") from exc + + @classmethod + def from_bytes(cls, data: bytes, fs_cache: Optional[dict] = None): + from dvc.fs import get_fs_cls + from dvc.fs.repo import RepoFileSystem + + try: + dict_ = pickle.loads(data) + except pickle.PickleError as exc: + raise ObjectFormatError("ReferenceHashFile is corrupted") from exc + + try: + path_info = dict_[cls.PARAM_PATH] + hash_info = dict_[cls.PARAM_HASH] + except KeyError as exc: + raise ObjectFormatError("ReferenceHashFile is corrupted") from exc + + scheme, config_pairs = dict_.get(cls.PARAM_FS_CONFIG) + fs = fs_cache.get((scheme, config_pairs)) if fs_cache else None + if not fs: + config = dict(config_pairs) + if RepoFileSystem.PARAM_REPO_URL in config: + fs = RepoFileSystem(**config) + path_info = fs.root_dir / path_info + else: + fs_cls = get_fs_cls(config, scheme=path_info.scheme) + fs = fs_cls(**config) + return ReferenceHashFile( + path_info, + fs, + hash_info, + mtime=dict_.get(cls.PARAM_MTIME), + size=dict_.get(cls.PARAM_SIZE), + ) + + @staticmethod + def config_tuple(fs: "BaseFileSystem"): + return ( + fs.scheme, + tuple( + (key, value) + for key, value in sorted( + fs.config.items(), key=lambda item: item[0] + ) + ), + ) diff --git a/dvc/objects/stage.py b/dvc/objects/stage.py index ea0a62254..c86fc1b48 100644 --- a/dvc/objects/stage.py +++ b/dvc/objects/stage.py @@ -1,6 +1,10 @@ +import hashlib +import logging +import os +import pathlib from concurrent.futures import ThreadPoolExecutor from functools import partial -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING, Dict, Tuple from dvc.exceptions import DvcIgnoreInCollectedDirError from dvc.hash_info import HashInfo @@ -8,32 +12,34 @@ from dvc.ignore import DvcIgnore from dvc.progress import Tqdm from dvc.utils import file_md5 +from .db.reference import ReferenceObjectDB from .file import HashFile if TYPE_CHECKING: from dvc.fs.base import BaseFileSystem - from dvc.types import DvcPath + from dvc.types import AnyPath, DvcPath from .db.base import ObjectDB - -_STAGING_DIR = "staging" -_STAGING_MEMFS_PATH = f"dvc-{_STAGING_DIR}" +logger = logging.getLogger(__name__) -def _upload_file(path_info, fs, odb): +_STAGING_MEMFS_PATH = "dvc-staging" + + +def _upload_file(path_info, fs, odb, upload_odb): from dvc.utils import tmp_fname from dvc.utils.stream import HashedStreamReader - tmp_info = odb.path_info / tmp_fname() + tmp_info = upload_odb.path_info / tmp_fname() with fs.open(path_info, mode="rb", chunk_size=fs.CHUNK_SIZE) as stream: stream = HashedStreamReader(stream) - odb.fs.upload_fobj( + upload_odb.fs.upload_fobj( stream, tmp_info, desc=path_info.name, size=fs.getsize(path_info) ) - obj = HashFile(tmp_info, odb.fs, stream.hash_info) - return path_info, obj + odb.add(tmp_info, upload_odb.fs, stream.hash_info) + return path_info, odb.get(stream.hash_info) def _get_file_hash(path_info, fs, name): @@ -71,15 +77,20 @@ def get_file_hash(path_info, fs, name, state=None): return hash_info -def _get_file_obj(path_info, fs, name, odb=None, upload=False): - if upload: - assert odb and name == "md5" - return _upload_file(path_info, fs, odb) - +def _get_file_obj( + path_info, fs, name, odb=None, upload_odb=None, dry_run=False +): state = odb.state if odb else None - obj = HashFile( - path_info, fs, get_file_hash(path_info, fs, name, state=state) - ) + hash_info = get_file_hash(path_info, fs, name, state=state) + if upload_odb and not dry_run: + assert odb and name == "md5" + return _upload_file(path_info, fs, odb, upload_odb) + + if dry_run: + obj = HashFile(path_info, fs, hash_info) + else: + odb.add(path_info, fs, hash_info, move=False) + obj = odb.get(hash_info) return path_info, obj @@ -87,9 +98,9 @@ def _build_objects( path_info, fs, name, - odb=None, - upload=False, dvcignore=None, + jobs=None, + no_progress_bar=False, **kwargs, ): if dvcignore: @@ -99,33 +110,24 @@ def _build_objects( with Tqdm( unit="md5", desc="Computing file/dir hashes (only done once)", - disable=kwargs.pop("no_progress_bar", False), + disable=no_progress_bar, ) as pbar: worker = pbar.wrap_fn( partial( _get_file_obj, fs=fs, name=name, - odb=odb, - upload=upload, + **kwargs, ) ) with ThreadPoolExecutor( - max_workers=kwargs.pop("jobs", fs.hash_jobs) + max_workers=jobs if jobs is not None else fs.hash_jobs ) as executor: yield from executor.map(worker, walk_iterator) -def _iter_objects(path_info, fs, name, upload=False, **kwargs): - if not upload and name in fs.DETAIL_FIELDS: - for details in fs.find(path_info, detail=True): - file_info = path_info.replace(path=details["name"]) - hash_info = HashInfo(name, details[name], size=details.get("size")) - yield file_info, HashFile(file_info, fs, hash_info) - - return None - - yield from _build_objects(path_info, fs, name, upload=upload, **kwargs) +def _iter_objects(path_info, fs, name, **kwargs): + yield from _build_objects(path_info, fs, name, **kwargs) def _build_tree(path_info, fs, name, **kwargs): @@ -145,14 +147,13 @@ def _build_tree(path_info, fs, name, **kwargs): # Yes, this is a BUG, as long as we permit "/" in # filenames on Windows and "\" on Unix tree.add(file_info.relative_to(path_info).parts, obj) - tree.digest() return tree -def _get_tree_obj(path_info, fs, name, odb=None, **kwargs): +def _get_tree_obj(path_info, fs, fs_info, name, odb=None, **kwargs): from .tree import Tree - value = fs.info(path_info).get(name) + value = fs_info.get(name) if odb and value: hash_info = HashInfo(name, value) try: @@ -168,40 +169,59 @@ def _get_tree_obj(path_info, fs, name, odb=None, **kwargs): pass tree = _build_tree(path_info, fs, name, odb=odb, **kwargs) + state = odb.state if odb and odb.state else None + hash_info = None + if state: + hash_info = state.get( # pylint: disable=assignment-from-none + path_info, fs + ) + tree.digest(hash_info=hash_info) + odb.add(tree.path_info, tree.fs, tree.hash_info, move=True) + raw = odb.get(tree.hash_info) + # cleanup unneeded memfs tmpfile and return tree based on the + # ODB fs/path + if odb.fs != tree.fs: + tree.fs.remove(tree.path_info) + tree.fs = raw.fs + tree.path_info = raw.path_info return tree -def get_staging(odb: Optional["ObjectDB"] = None) -> "ObjectDB": +_url_cache: Dict[str, str] = {} + + +def _make_staging_url(path_info: "AnyPath"): + from dvc.path_info import CloudURLInfo + from dvc.scheme import Schemes + + url = CloudURLInfo(f"{Schemes.MEMORY}://{_STAGING_MEMFS_PATH}") + if path_info: + if isinstance(path_info, (str, pathlib.PurePath)): + path = os.path.abspath(path_info) + else: + path = str(path_info) + if path not in _url_cache: + _url_cache[path] = hashlib.sha256(path.encode("utf-8")).hexdigest() + url /= _url_cache[path] + return url + + +def _get_staging(odb: "ObjectDB") -> "ObjectDB": """Return an ODB that can be used for staging objects. - If odb.fs is local, .dvc/tmp/staging will be returned. Otherwise - the the global (temporary) memfs ODB will be returned. + Staging will be a reference ODB stored in the the global memfs. """ from dvc.fs.memory import MemoryFileSystem - from dvc.path_info import CloudURLInfo, PathInfo - from dvc.scheme import Schemes - from .db import get_odb - - if odb and odb.fs.scheme == Schemes.LOCAL and odb.tmp_dir: - fs = odb.fs - path_info: "DvcPath" = PathInfo(odb.tmp_dir) / _STAGING_DIR - config = odb.config - else: - fs = MemoryFileSystem() - path_info = CloudURLInfo(f"{Schemes.MEMORY}://{_STAGING_MEMFS_PATH}") - config = {} - return get_odb(fs, path_info, **config) - - -def is_memfs_staging(odb: "ObjectDB"): - staging = get_staging() - return odb.fs == staging.fs and odb.path_info == staging.path_info + fs = MemoryFileSystem() + path_info = _make_staging_url(odb.path_info) + state = odb.state + return ReferenceObjectDB(fs, path_info, state=state) def _load_from_state(odb, staging, path_info, fs, name): - from . import load + from . import check, load from .errors import ObjectFormatError from .tree import Tree @@ -212,6 +232,7 @@ def _load_from_state(odb, staging, path_info, fs, name): if odb_.exists(hash_info): try: obj = load(odb_, hash_info) + check(odb_, obj, check_hash=False) if isinstance(obj, Tree): obj.hash_info.nfiles = len(obj) for key, entry in obj: @@ -222,8 +243,8 @@ def _load_from_state(odb, staging, path_info, fs, name): obj.path_info = path_info assert obj.hash_info.name == name obj.hash_info.size = hash_info.size - return obj - except ObjectFormatError: + return odb_, obj + except (ObjectFormatError, FileNotFoundError): pass raise FileNotFoundError @@ -248,17 +269,31 @@ def _stage_external_tree_info(odb, tree, name): def stage( - odb: Optional["ObjectDB"], + odb: "ObjectDB", path_info: "DvcPath", fs: "BaseFileSystem", name: str, upload: bool = False, + dry_run: bool = False, **kwargs, -) -> "HashFile": +) -> Tuple["ObjectDB", "HashFile"]: + """Stage (prepare) objects from the given path for addition to an ODB. + + Returns at tuple of (staging_odb, object) where addition to the ODB can + be completed by transferring the object from staging to the dest ODB. + + If dry_run is True, object hashes will be computed and returned, but file + objects themselves will not be added to the staging ODB (i.e. the resulting + file objects cannot transferred from staging to another ODB). + + If upload is True, files will be uploaded to a temporary path on the dest + ODB filesystem, and staged objects will reference the uploaded path rather + than the original source path. + """ assert path_info and path_info.scheme == fs.scheme details = fs.info(path_info) - staging = get_staging(odb) + staging = _get_staging(odb) if odb: try: return _load_from_state(odb, staging, path_info, fs, name) @@ -267,26 +302,29 @@ def stage( if details["type"] == "directory": obj = _get_tree_obj( - path_info, fs, name, odb=odb, upload=upload, **kwargs + path_info, + fs, + details, + name, + odb=staging, + upload_odb=odb if upload else None, + dry_run=dry_run, + **kwargs, ) - if name == "md5": - if odb and odb.exists(obj.hash_info): - raw = odb.get(obj.hash_info) - else: - staging.add(obj.path_info, obj.fs, obj.hash_info) - raw = staging.get(obj.hash_info) - # cleanup unneeded memfs tmpfile and return obj based on staging - # ODB fs/path - if obj.fs != raw.fs: - obj.fs.remove(obj.path_info) - obj.fs = raw.fs - obj.path_info = raw.path_info - else: + logger.debug("staged tree '%s'", obj) + if name != "md5": obj = _stage_external_tree_info(odb, obj, name) else: - _, obj = _get_file_obj(path_info, fs, name, odb=odb, upload=upload) + _, obj = _get_file_obj( + path_info, + fs, + name, + odb=staging, + upload_odb=odb if upload else None, + dry_run=dry_run, + ) if odb and odb.state and obj.hash_info: odb.state.save(path_info, fs, obj.hash_info) - return obj + return staging, obj diff --git a/dvc/objects/status.py b/dvc/objects/status.py index ca2fc4cd3..0dfaabea1 100644 --- a/dvc/objects/status.py +++ b/dvc/objects/status.py @@ -2,6 +2,7 @@ import logging from typing import TYPE_CHECKING, Dict, Iterable, NamedTuple, Optional, Set from dvc.hash_info import HashInfo +from dvc.scheme import Schemes from .tree import Tree @@ -14,23 +15,21 @@ logger = logging.getLogger(__name__) class StatusResult(NamedTuple): - exists: Set["HashFile"] - missing: Set["HashFile"] + exists: Set["HashInfo"] + missing: Set["HashInfo"] class CompareStatusResult(NamedTuple): - ok: Set["HashFile"] - missing: Set["HashFile"] - new: Set["HashFile"] - deleted: Set["HashFile"] + ok: Set["HashInfo"] + missing: Set["HashInfo"] + new: Set["HashInfo"] + deleted: Set["HashInfo"] -def _indexed_dir_hashes(odb, index, hash_infos, name, cache_odb): - from . import load - +def _indexed_dir_hashes(odb, index, dir_objs, name, cache_odb): # Validate our index by verifying all indexed .dir hashes # still exist on the remote - dir_hashes = {hash_info.value for hash_info in hash_infos} + dir_hashes = set(dir_objs.keys()) indexed_dirs = set(index.dir_hashes()) indexed_dir_exists = set() if indexed_dirs: @@ -51,10 +50,12 @@ def _indexed_dir_hashes(odb, index, hash_infos, name, cache_odb): # If .dir hash exists in the ODB, assume directory contents # also exists for dir_hash in dir_exists: - try: - tree = load(cache_odb, HashInfo(name, dir_hash)) - except FileNotFoundError: - continue + tree = dir_objs.get(dir_hash) + if not tree: + try: + tree = Tree.load(cache_odb, HashInfo(name, dir_hash)) + except FileNotFoundError: + continue file_hashes = [entry.hash_info.value for _, entry in tree] if dir_hash not in index: logger.debug( @@ -67,21 +68,13 @@ def _indexed_dir_hashes(odb, index, hash_infos, name, cache_odb): yield tree.hash_info.value -def _status_staging(objs: Iterable["HashFile"]) -> "StatusResult": - exists: Set["HashFile"] = set() - for obj in objs: - if isinstance(obj, Tree): - exists.update(entry for _, entry in obj) - exists.add(obj) - return StatusResult(exists, set()) - - def status( odb: "ObjectDB", - objs: Iterable["HashFile"], + obj_ids: Iterable["HashInfo"], name: Optional[str] = None, index: Optional["ObjectDBIndexBase"] = None, cache_odb: Optional["ObjectDB"] = None, + shallow: bool = True, **kwargs, ) -> "StatusResult": """Return status of whether or not the specified objects exist odb. @@ -93,39 +86,41 @@ def status( exists: objs that exist in odb missing: objs that do not exist in ODB """ - from .stage import is_memfs_staging - logger.debug("Preparing to collect status from '%s'", odb.path_info) if not name: name = odb.fs.PARAM_CHECKSUM - if is_memfs_staging(odb): - # assume memfs staged objects already exist - return _status_staging(objs) + if cache_odb is None: + cache_odb = odb - hash_objs: Dict[str, "HashFile"] = {} - hashes: Set[str] = set() - dir_infos: Set["HashInfo"] = set() - exists: Set[str] = set() - for obj in objs: - assert obj.hash_info and obj.hash_info.value - if isinstance(obj, Tree): - for _, entry in obj: - assert entry.hash_info and entry.hash_info.value - hash_objs[entry.hash_info.value] = entry - hashes.add(entry.hash_info.value) + hash_infos: Dict[str, "HashInfo"] = {} + dir_objs: Dict[str, Optional["HashFile"]] = {} + for hash_info in obj_ids: + assert hash_info.value + if hash_info.isdir: + if shallow: + tree = None + else: + tree = Tree.load(cache_odb, hash_info) + for _, entry in tree: + assert entry.hash_info and entry.hash_info.value + hash_infos[entry.hash_info.value] = entry.hash_info if index: - dir_infos.add(obj.hash_info) - hashes.add(obj.hash_info.value) - hash_objs[obj.hash_info.value] = obj + dir_objs[hash_info.value] = tree + hash_infos[hash_info.value] = hash_info + + if odb.fs.scheme == Schemes.MEMORY: + # assume memfs staged objects already exist + return StatusResult(set(hash_infos.values()), set()) + + hashes: Set[str] = set(hash_infos.keys()) + exists: Set[str] = set() logger.debug("Collecting status from '%s'", odb.path_info) if index and hashes: - if dir_infos: - if cache_odb is None: - cache_odb = odb + if dir_objs: exists = hashes.intersection( - _indexed_dir_hashes(odb, index, dir_infos, name, cache_odb) + _indexed_dir_hashes(odb, index, dir_objs, name, cache_odb) ) hashes.difference_update(exists) if hashes: @@ -137,15 +132,15 @@ def status( odb.hashes_exist(hashes, name=str(odb.path_info), **kwargs) ) return StatusResult( - {hash_objs[hash_] for hash_ in exists}, - {hash_objs[hash_] for hash_ in (hashes - exists)}, + {hash_infos[hash_] for hash_ in exists}, + {hash_infos[hash_] for hash_ in (hashes - exists)}, ) def compare_status( src: "ObjectDB", dest: "ObjectDB", - objs: Iterable["HashFile"], + obj_ids: Iterable["HashInfo"], log_missing: bool = True, check_deleted: bool = True, src_index: Optional["ObjectDBIndexBase"] = None, @@ -160,11 +155,17 @@ def compare_status( new: hashes that only exist in src deleted: hashes that only exist in dest """ - dest_exists, dest_missing = status(dest, objs, index=dest_index, **kwargs) + if "cache_odb" not in kwargs: + kwargs["cache_odb"] = src + dest_exists, dest_missing = status( + dest, obj_ids, index=dest_index, **kwargs + ) # for transfer operations we can skip src status check when all objects # already exist in dest if dest_missing or check_deleted: - src_exists, src_missing = status(src, objs, index=src_index, **kwargs) + src_exists, src_missing = status( + src, obj_ids, index=src_index, **kwargs + ) else: src_exists = dest_exists src_missing = set() @@ -176,7 +177,8 @@ def compare_status( ) if log_missing and result.missing: missing_desc = "\n".join( - f"name: {obj.name}, {obj.hash_info}" for obj in result.missing + f"name: {hash_info.obj_name}, {hash_info}" + for hash_info in result.missing ) logger.warning( "Some of the cache files do not exist neither locally " diff --git a/dvc/objects/transfer.py b/dvc/objects/transfer.py index b59f1b568..3c2b151eb 100644 --- a/dvc/objects/transfer.py +++ b/dvc/objects/transfer.py @@ -4,14 +4,15 @@ import logging from concurrent import futures from concurrent.futures import ThreadPoolExecutor from functools import partial, wraps -from typing import TYPE_CHECKING, Iterable, Optional, Set +from typing import TYPE_CHECKING, Callable, Iterable, Optional + +from funcy import split from dvc.progress import Tqdm -from .file import HashFile -from .tree import Tree - if TYPE_CHECKING: + from dvc.hash_info import HashInfo + from .db.base import ObjectDB from .db.index import ObjectDBIndexBase @@ -20,9 +21,9 @@ logger = logging.getLogger(__name__) def _log_exceptions(func): @wraps(func) - def wrapper(odb, obj, *args, **kwargs): + def wrapper(path_info, *args, **kwargs): try: - func(odb, obj, *args, **kwargs) + func(path_info, *args, **kwargs) return 0 except Exception as exc: # pylint: disable=broad-except # NOTE: this means we ran out of file descriptors and there is no @@ -31,23 +32,17 @@ def _log_exceptions(func): if isinstance(exc, OSError) and exc.errno == errno.EMFILE: raise - logger.exception( - "failed to transfer '%s' to '%s'", - obj.path_info, - odb.get(obj.hash_info).path_info, - ) + logger.exception("failed to transfer '%s'", path_info) return 1 return wrapper -def _transfer(src, dest, dir_objs, file_objs, missing, jobs, verify, **kwargs): - from . import save - from .stage import is_memfs_staging - - is_staged = is_memfs_staging(src) - func = _log_exceptions(save) - total = len(dir_objs) + len(file_objs) +def _transfer( + src, dest, dir_ids, file_ids, missing_ids, jobs, verify, move, **kwargs +): + func = _log_exceptions(dest.add) + total = len(dir_ids) + len(file_ids) if total == 0: return 0 with Tqdm(total=total, unit="file", desc="Transferring") as pbar: @@ -59,37 +54,43 @@ def _transfer(src, dest, dir_objs, file_objs, missing, jobs, verify, **kwargs): jobs, func, src, - dest, - is_staged, verify, + move, ) - processor.save_func = func + processor.add_func = func _do_transfer( src, dest, - dir_objs, - file_objs, - {obj.hash_info for obj in missing}, + dir_ids, + file_ids, + missing_ids, processor, + verify=verify, + move=move, **kwargs, ) return total -def _create_tasks(executor, jobs, func, src, dest, is_staged, verify, objs): +def _create_tasks(executor, jobs, func, src, verify, move, obj_ids): fails = 0 - obj_iter = iter(objs) + hash_iter = iter(obj_ids) + + def submit(hash_info): + obj = src.get(hash_info) + return executor.submit( + func, + obj.path_info, + obj.fs, + obj.hash_info, + verify=verify, + move=move, + ) def create_taskset(amount): return { - executor.submit( - func, - dest, - _raw_obj(src, obj, is_staged), - move=False, - verify=verify, - ) - for obj in itertools.islice(obj_iter, amount) + submit(hash_info) + for hash_info in itertools.islice(hash_iter, amount) } tasks = create_taskset(jobs * 5) @@ -100,39 +101,46 @@ def _create_tasks(executor, jobs, func, src, dest, is_staged, verify, objs): return fails -def _raw_obj(odb, obj, is_staged=False): - if is_staged: - return HashFile(obj.path_info, obj.fs, obj.hash_info) - return odb.get(obj.hash_info) - - def _do_transfer( - src, - dest, - dir_objs, - file_objs, - missing_hashes, - processor, + src: "ObjectDB", + dest: "ObjectDB", + dir_ids: Iterable["HashInfo"], + file_ids: Iterable["HashInfo"], + missing_ids: Iterable["HashInfo"], + processor: Callable, src_index: Optional["ObjectDBIndexBase"] = None, dest_index: Optional["ObjectDBIndexBase"] = None, + cache_odb: Optional["ObjectDB"] = None, **kwargs, ): from dvc.exceptions import FileTransferError + from dvc.objects.errors import ObjectFormatError total_fails = 0 succeeded_dir_objs = [] - all_file_objs = set(file_objs) + all_file_ids = set(file_ids) - for dir_obj in dir_objs: - bound_file_objs = set() - directory_hashes = {entry.hash_info for _, entry in dir_obj} + for dir_hash in dir_ids: + from .tree import Tree - for file_obj in all_file_objs.copy(): - if file_obj.hash_info in directory_hashes: - bound_file_objs.add(file_obj) - all_file_objs.remove(file_obj) + bound_file_ids = set() + dir_obj: Optional["Tree"] = None + for odb in (cache_odb, src): + if odb is not None: + try: + dir_obj = Tree.load(odb, dir_hash) + break + except (FileNotFoundError, ObjectFormatError): + pass + assert dir_obj + entry_ids = {entry.hash_info for _, entry in dir_obj} - dir_fails = processor(bound_file_objs) + for file_hash in all_file_ids.copy(): + if file_hash in entry_ids: + bound_file_ids.add(file_hash) + all_file_ids.remove(file_hash) + + dir_fails = processor(bound_file_ids) if dir_fails: logger.debug( "failed to upload full contents of '%s', " @@ -145,7 +153,7 @@ def _do_transfer( dest.get(dir_obj.hash_info).path_info, ) total_fails += dir_fails + 1 - elif directory_hashes.intersection(missing_hashes): + elif entry_ids.intersection(missing_ids): # if for some reason a file contained in this dir is # missing both locally and in the remote, we want to # push whatever file content we have, but should not @@ -156,15 +164,19 @@ def _do_transfer( dir_obj.name, ) else: - is_dir_failed = processor.save_func( - dest, src.get(dir_obj.hash_info), move=False + raw_obj = src.get(dir_obj.hash_info) + is_dir_failed = processor.add_func( # type: ignore[attr-defined] + raw_obj.path_info, + raw_obj.fs, + raw_obj.hash_info, + **kwargs, ) total_fails += is_dir_failed if not is_dir_failed: succeeded_dir_objs.append(dir_obj) # insert the rest - total_fails += processor(all_file_objs) + total_fails += processor(all_file_ids) if total_fails: if src_index: src_index.clear() @@ -179,15 +191,17 @@ def _do_transfer( dir_obj.hash_info, len(file_hashes), ) + assert dir_obj.hash_info and dir_obj.hash_info.value dest_index.update([dir_obj.hash_info.value], file_hashes) def transfer( src: "ObjectDB", dest: "ObjectDB", - objs: Iterable["HashFile"], + obj_ids: Iterable["HashInfo"], jobs: Optional[int] = None, verify: bool = False, + move: bool = False, **kwargs, ) -> int: """Transfer (copy) the specified objects from one ODB to another. @@ -201,27 +215,25 @@ def transfer( src.path_info, dest.path_info, ) - status = compare_status(src, dest, objs, check_deleted=False, **kwargs) + if src == dest: + return 0 + + status = compare_status(src, dest, obj_ids, check_deleted=False, **kwargs) if not status.new: return 0 - files: Set["HashFile"] = set() - dirs: Set["HashFile"] = set() - for obj in status.new: - if isinstance(obj, Tree): - dirs.add(obj) - else: - files.add(obj) + dir_ids, file_ids = split(lambda hash_info: hash_info.isdir, status.new) if jobs is None: jobs = dest.fs.jobs return _transfer( src, dest, - dirs, - files, + set(dir_ids), + set(file_ids), status.missing, jobs, verify, + move, **kwargs, ) diff --git a/dvc/objects/tree.py b/dvc/objects/tree.py index cc2c4dd53..fad5c0439 100644 --- a/dvc/objects/tree.py +++ b/dvc/objects/tree.py @@ -1,7 +1,7 @@ import json import logging import posixpath -from typing import Dict, Optional, Tuple +from typing import TYPE_CHECKING, Dict, Optional, Tuple from funcy import cached_property @@ -9,6 +9,9 @@ from .errors import ObjectFormatError from .file import HashFile from .stage import get_file_hash +if TYPE_CHECKING: + from dvc.hash_info import HashInfo + logger = logging.getLogger(__name__) @@ -29,7 +32,7 @@ class Tree(HashFile): self.__dict__.pop("trie", None) self._dict[key] = obj - def digest(self): + def digest(self, hash_info: Optional["HashInfo"] = None): from dvc.fs.memory import MemoryFileSystem from dvc.path_info import CloudURLInfo from dvc.utils import tmp_fname @@ -40,8 +43,12 @@ class Tree(HashFile): fobj.write(self.as_bytes()) self.fs = memfs self.path_info = path_info - self.hash_info = get_file_hash(path_info, memfs, "md5") - self.hash_info.value += ".dir" + if hash_info: + self.hash_info = hash_info + else: + self.hash_info = get_file_hash(path_info, memfs, "md5") + assert self.hash_info.value + self.hash_info.value += ".dir" try: self.hash_info.size = sum(obj.size for _, obj in self) except TypeError: diff --git a/dvc/output.py b/dvc/output.py index 5ed838021..d3d3bf343 100644 --- a/dvc/output.py +++ b/dvc/output.py @@ -25,16 +25,15 @@ from .fs.s3 import S3FileSystem from .hash_info import HashInfo from .istextfile import istextfile from .objects import Tree -from .objects import save as osave from .objects.errors import ObjectFormatError from .objects.stage import stage as ostage +from .objects.transfer import transfer as otransfer from .scheme import Schemes from .utils import relpath from .utils.fs import path_isin if TYPE_CHECKING: from .objects.db.base import ObjectDB - from .objects.file import HashFile logger = logging.getLogger(__name__) @@ -404,21 +403,21 @@ class Output: return self.odb.hash_to_path_info(self.hash_info.value).url def get_hash(self): - if not self.use_cache: - return ostage( - self.repo.odb.local, - self.path_info, - self.fs, - self.fs.PARAM_CHECKSUM, - dvcignore=self.dvcignore, - ).hash_info - return ostage( - self.odb, + if self.use_cache: + odb = self.odb + name = self.odb.fs.PARAM_CHECKSUM + else: + odb = self.repo.odb.local + name = self.fs.PARAM_CHECKSUM + _, obj = ostage( + odb, self.path_info, self.fs, - self.odb.fs.PARAM_CHECKSUM, + name, dvcignore=self.dvcignore, - ).hash_info + dry_run=not self.use_cache, + ) + return obj.hash_info @property def is_dir_checksum(self): @@ -545,7 +544,7 @@ class Output: logger.debug("Output '%s' didn't change. Skipping saving.", self) return - self.obj = ostage( + _, self.obj = ostage( self.odb, self.path_info, self.fs, @@ -574,14 +573,20 @@ class Output: if granular: obj = self._commit_granular_dir(filter_info) else: - obj = ostage( + staging, obj = ostage( self.odb, filter_info or self.path_info, self.fs, self.odb.fs.PARAM_CHECKSUM, dvcignore=self.dvcignore, ) - osave(self.odb, obj) + otransfer( + staging, + self.odb, + {obj.hash_info}, + shallow=False, + move=True, + ) checkout( filter_info or self.path_info, self.fs, @@ -595,7 +600,7 @@ class Output: def _commit_granular_dir(self, filter_info): prefix = filter_info.relative_to(self.path_info).parts - save_obj = ostage( + staging, save_obj = ostage( self.odb, self.path_info, self.fs, @@ -604,7 +609,13 @@ class Output: ) save_obj = save_obj.filter(prefix) checkout_obj = save_obj.get(prefix) - osave(self.odb, save_obj) + otransfer( + staging, + self.odb, + {save_obj.hash_info} | {entry.hash_info for _, entry in save_obj}, + shallow=True, + move=True, + ) return checkout_obj def dumpd(self): @@ -786,7 +797,7 @@ class Output: upload = not (update and from_fs.isdir(from_info)) jobs = jobs or min((from_fs.jobs, odb.fs.jobs)) - obj = ostage( + staging, obj = ostage( odb, from_info, from_fs, @@ -795,7 +806,14 @@ class Output: jobs=jobs, no_progress_bar=no_progress_bar, ) - osave(odb, obj, jobs=jobs, move=upload) + otransfer( + staging, + odb, + {obj.hash_info}, + jobs=jobs, + move=upload, + shallow=False, + ) self.hash_info = obj.hash_info return obj @@ -825,7 +843,7 @@ class Output: try: objects.check(self.odb, obj) except FileNotFoundError: - self.repo.cloud.pull([obj], **kwargs) + self.repo.cloud.pull([obj.hash_info], **kwargs) if self.obj: return self.obj @@ -839,8 +857,8 @@ class Output: def collect_used_dir_cache( self, remote=None, force=False, jobs=None, filter_info=None - ) -> Dict[Optional["ObjectDB"], Set["HashFile"]]: - """Fetch dir cache and return used objects for this out.""" + ) -> Dict[Optional["ObjectDB"], Set["HashInfo"]]: + """Fetch dir cache and return used object IDs for this out.""" try: self.get_dir_cache(jobs=jobs, remote=remote) @@ -866,13 +884,12 @@ class Output: if filter_info and filter_info != self.path_info: prefix = filter_info.relative_to(self.path_info).parts obj = obj.filter(prefix) - self._set_obj_names(obj) - return {None: {obj}} + return {None: set(self._named_obj_ids(obj))} def get_used_objs( self, **kwargs - ) -> Dict[Optional["ObjectDB"], Set["HashFile"]]: - """Return filtered set of used objects for this out.""" + ) -> Dict[Optional["ObjectDB"], Set["HashInfo"]]: + """Return filtered set of used object IDs for this out.""" if not self.use_cache: return {} @@ -905,19 +922,21 @@ class Output: obj = self.get_obj(filter_info=kwargs.get("filter_info")) if not obj: obj = self.odb.get(self.hash_info) - self._set_obj_names(obj) - return {None: {obj}} + return {None: set(self._named_obj_ids(obj))} - def _set_obj_names(self, obj): - obj.name = str(self) + def _named_obj_ids(self, obj): + name = str(self) + obj.hash_info.obj_name = name + yield obj.hash_info if isinstance(obj, Tree): for key, entry_obj in obj: - entry_obj.name = self.fs.sep.join([obj.name, *key]) + entry_obj.hash_info.obj_name = self.fs.sep.join([name, *key]) + yield entry_obj.hash_info def get_used_external( self, **kwargs - ) -> Dict[Optional["ObjectDB"], Set["HashFile"]]: + ) -> Dict[Optional["ObjectDB"], Set["HashInfo"]]: if not self.use_cache or not self.stage.is_repo_import: return {} diff --git a/dvc/path_info.py b/dvc/path_info.py index 0ed5437e2..497d606c7 100644 --- a/dvc/path_info.py +++ b/dvc/path_info.py @@ -64,6 +64,8 @@ class PathInfo(pathlib.PurePath, _BasePath): url = fspath + path = fspath + def relpath(self, other): return self.__class__(relpath(self, other)) diff --git a/dvc/repo/diff.py b/dvc/repo/diff.py index 74431c3b0..fbbd32aa3 100644 --- a/dvc/repo/diff.py +++ b/dvc/repo/diff.py @@ -134,8 +134,12 @@ def _output_paths(repo, repo_fs, targets): def _to_checksum(output): if on_working_fs: return ostage( - repo.odb.local, output.path_info, repo.odb.local.fs, "md5" - ).hash_info.value + repo.odb.local, + output.path_info, + repo.odb.local.fs, + "md5", + dry_run=True, + )[1].hash_info.value return output.hash_info.value for stage in repo.stages: diff --git a/dvc/repo/fetch.py b/dvc/repo/fetch.py index ff9d0f5cd..d9cee65aa 100644 --- a/dvc/repo/fetch.py +++ b/dvc/repo/fetch.py @@ -1,9 +1,8 @@ import logging -from dvc.config import NoRemoteError from dvc.exceptions import DownloadError, FileTransferError +from dvc.scheme import Schemes -from ..scheme import Schemes from . import locked logger = logging.getLogger(__name__) @@ -60,10 +59,13 @@ def fetch( except DownloadError as exc: failed += exc.amount - for odb, objs in used.items(): + for odb, obj_ids in sorted( + used.items(), + key=lambda item: item[0] and item[0].fs.scheme == Schemes.MEMORY, + ): d, f = _fetch( self, - objs, + obj_ids, jobs=jobs, remote=remote, odb=odb, @@ -77,14 +79,11 @@ def fetch( return downloaded -def _fetch(repo, objs, **kwargs): +def _fetch(repo, obj_ids, **kwargs): downloaded = 0 failed = 0 try: - downloaded += repo.cloud.pull(objs, **kwargs) - except NoRemoteError: - if any(obj.fs.scheme == Schemes.LOCAL for obj in objs): - raise + downloaded += repo.cloud.pull(obj_ids, **kwargs) except FileTransferError as exc: failed += exc.amount return downloaded, failed diff --git a/dvc/repo/gc.py b/dvc/repo/gc.py index ceb0a5631..08717658b 100644 --- a/dvc/repo/gc.py +++ b/dvc/repo/gc.py @@ -1,14 +1,9 @@ import logging -from typing import TYPE_CHECKING, Set from dvc.exceptions import InvalidArgumentError -from ..scheme import Schemes from . import locked -if TYPE_CHECKING: - from dvc.objects.file import HashFile - logger = logging.getLogger(__name__) @@ -50,20 +45,19 @@ def gc( from contextlib import ExitStack from dvc.objects.db import get_index - from dvc.objects.stage import get_staging from dvc.repo import Repo if not repos: repos = [] all_repos = [Repo(path) for path in repos] - used_objs: Set["HashFile"] = set() + used_obj_ids = set() with ExitStack() as stack: for repo in all_repos: stack.enter_context(repo.lock) for repo in all_repos + [self]: - for objs in repo.used_objs( + for obj_ids in repo.used_objs( all_branches=all_branches, with_deps=with_deps, all_tags=all_tags, @@ -73,28 +67,21 @@ def gc( force=force, jobs=jobs, ).values(): - used_objs.update(objs) + used_obj_ids.update(obj_ids) for scheme, odb in self.odb.by_scheme(): if not odb: continue - removed = odb.gc( - {obj for obj in used_objs if obj.fs.scheme == scheme}, - jobs=jobs, - ) + removed = odb.gc(used_obj_ids, jobs=jobs) if not removed: logger.info(f"No unused '{scheme}' cache to remove.") - get_staging(self.odb.local).gc(set(), jobs=jobs) if not cloud: return odb = self.cloud.get_remote_odb(remote, "gc -c") - removed = odb.gc( - {obj for obj in used_objs if obj.fs.scheme == Schemes.LOCAL}, - jobs=jobs, - ) + removed = odb.gc(used_obj_ids) if removed: get_index(odb).clear() else: diff --git a/dvc/repo/push.py b/dvc/repo/push.py index 341e0a720..88f9adc6a 100644 --- a/dvc/repo/push.py +++ b/dvc/repo/push.py @@ -41,11 +41,11 @@ def push( ) pushed = len(used_run_cache) - for odb, objs in used.items(): + for odb, obj_ids in used.items(): if odb and odb.read_only: continue try: - pushed += self.cloud.push(objs, jobs, remote=remote, odb=odb) + pushed += self.cloud.push(obj_ids, jobs, remote=remote, odb=odb) except FileTransferError as exc: raise UploadError(exc.amount) return pushed diff --git a/dvc/repo/status.py b/dvc/repo/status.py index db7371756..2d0edb6f7 100644 --- a/dvc/repo/status.py +++ b/dvc/repo/status.py @@ -89,16 +89,16 @@ def _cloud_status( ) ret = {} - for odb, objs in used.items(): + for odb, obj_ids in used.items(): if odb is not None: # ignore imported objects continue status_info = self.cloud.status( - objs, jobs, remote=remote, log_missing=False + obj_ids, jobs, remote=remote, log_missing=False ) for status_ in ("deleted", "new", "missing"): - for obj in getattr(status_info, status_, []): - ret[obj.name] = status_ + for hash_info in getattr(status_info, status_, []): + ret[hash_info.obj_name] = status_ return ret diff --git a/dvc/stage/__init__.py b/dvc/stage/__init__.py index c45d7df6c..7aa8a121c 100644 --- a/dvc/stage/__init__.py +++ b/dvc/stage/__init__.py @@ -36,8 +36,8 @@ from .utils import ( if TYPE_CHECKING: from dvc.dvcfile import DVCFile + from dvc.hash_info import HashInfo from dvc.objects.db.base import ObjectDB - from dvc.objects.file import HashFile logger = logging.getLogger(__name__) # Disallow all punctuation characters except hyphen and underscore @@ -639,8 +639,8 @@ class Stage(params.StageParams): def get_used_objs( self, *args, **kwargs - ) -> Dict[Optional["ObjectDB"], Set["HashFile"]]: - """Return set of objects used by this stage.""" + ) -> Dict[Optional["ObjectDB"], Set["HashInfo"]]: + """Return set of object IDs used by this stage.""" used_objs = defaultdict(set) for out in self.filter_outs(kwargs.get("filter_info")): for odb, objs in out.get_used_objs(*args, **kwargs).items(): diff --git a/dvc/state.py b/dvc/state.py index 65dca8d57..5b2f55d0a 100644 --- a/dvc/state.py +++ b/dvc/state.py @@ -89,10 +89,14 @@ class State(StateBase): # pylint: disable=too-many-instance-attributes inode = get_inode(path_info) logger.debug( - "state save (%s, %s, %s) %s", inode, mtime, size, hash_info.value + "state save (%s, %s, %s) %s", + inode, + mtime, + str(size), + hash_info.value, ) - self.md5s[inode] = (mtime, size, hash_info.value) + self.md5s[inode] = (mtime, str(size), hash_info.value) def get(self, path_info, fs): """Gets the hash for the specified path info. Hash will be @@ -122,10 +126,10 @@ class State(StateBase): # pylint: disable=too-many-instance-attributes value = self.md5s.get(inode) - if not value or value[0] != mtime or value[1] != size: + if not value or value[0] != mtime or value[1] != str(size): return None - return HashInfo("md5", value[2], size=int(size)) + return HashInfo("md5", value[2], size=size) def save_link(self, path_info, fs): """Adds the specified path to the list of links created by dvc. This diff --git a/dvc/utils/fs.py b/dvc/utils/fs.py index 33f49a358..4b80923e8 100644 --- a/dvc/utils/fs.py +++ b/dvc/utils/fs.py @@ -65,9 +65,7 @@ def get_mtime_and_size(path, fs, dvcignore=None): mtime = base_stat.st_mtime mtime = int(nanotime.timestamp(mtime)) - # State of files handled by dvc is stored in db as TEXT. - # We cast results to string for later comparisons with stored values. - return str(mtime), str(size) + return str(mtime), size class BasePathNotInCheckedPathException(DvcException): diff --git a/tests/func/objects/db/test_index.py b/tests/func/objects/db/test_index.py index 7e624776f..37dea8506 100644 --- a/tests/func/objects/db/test_index.py +++ b/tests/func/objects/db/test_index.py @@ -4,6 +4,7 @@ from dvc.exceptions import DownloadError, UploadError from dvc.fs.local import LocalFileSystem from dvc.objects.db import get_index from dvc.utils.fs import remove +from tests.utils import clean_staging @pytest.fixture @@ -16,6 +17,7 @@ def test_indexed_on_status(tmp_dir, dvc, index): foo = tmp_dir.dvc_gen({"foo": "foo content"})[0].outs[0] bar = tmp_dir.dvc_gen({"bar": {"baz": "baz content"}})[0].outs[0] baz_hash = bar.obj.trie.get(("baz",)).hash_info + clean_staging() dvc.push() index.clear() @@ -29,6 +31,7 @@ def test_indexed_on_push(tmp_dir, dvc, index): foo = tmp_dir.dvc_gen({"foo": "foo content"})[0].outs[0] bar = tmp_dir.dvc_gen({"bar": {"baz": "baz content"}})[0].outs[0] baz_hash = bar.obj.trie.get(("baz",)).hash_info + clean_staging() dvc.push() assert {bar.hash_info.value, baz_hash.value} == set(index.hashes()) diff --git a/tests/func/test_commit.py b/tests/func/test_commit.py index 21c231053..b507a552f 100644 --- a/tests/func/test_commit.py +++ b/tests/func/test_commit.py @@ -7,6 +7,7 @@ from dvc.dvcfile import PIPELINE_FILE from dvc.output import OutputDoesNotExistError from dvc.stage.exceptions import StageCommitError from dvc.utils.serialize import dump_yaml, load_yaml +from tests.utils import clean_staging def test_commit_recursive(tmp_dir, dvc): @@ -53,6 +54,8 @@ def test_commit_with_deps(tmp_dir, dvc, run_copy, run_kw): assert foo_stage.outs[0].changed_cache() assert stage.outs[0].changed_cache() + clean_staging() + dvc.commit(stage.path, with_deps=True) assert not foo_stage.outs[0].changed_cache() assert not stage.outs[0].changed_cache() @@ -66,6 +69,8 @@ def test_commit_changed_md5(tmp_dir, dvc): stage_file_content["md5"] = "1111111111" dump_yaml(stage.path, stage_file_content) + clean_staging() + with pytest.raises(StageCommitError): dvc.commit(stage.path) @@ -78,6 +83,8 @@ def test_commit_no_exec(tmp_dir, dvc): stage = dvc.run( name="my", cmd="mycmd", deps=["dep"], outs=["out"], no_exec=True ) + clean_staging() + assert dvc.status(stage.path) dvc.commit(stage.path, force=True) assert dvc.status(stage.path) == {} @@ -90,6 +97,7 @@ def test_commit_granular_output(tmp_dir, dvc): outs=["foo", "bar"], no_commit=True, ) + clean_staging() cache = tmp_dir / ".dvc" / "cache" assert not list(cache.glob("*/*")) @@ -103,6 +111,7 @@ def test_commit_granular_output(tmp_dir, dvc): def test_commit_granular_output_file(tmp_dir, dvc): tmp_dir.gen("foo", "foo") dvc.add("foo", no_commit=True) + clean_staging() dvc.commit("foo") assert dvc.status() == {} @@ -118,13 +127,12 @@ def test_commit_granular_output_dir(tmp_dir, dvc): } ) dvc.add("data", no_commit=True) + clean_staging() dvc.commit("data") assert dvc.status() == {} def test_commit_granular_dir(tmp_dir, dvc): - from dvc.objects.stage import _STAGING_DIR - tmp_dir.gen( { "data": { @@ -135,21 +143,18 @@ def test_commit_granular_dir(tmp_dir, dvc): } ) dvc.add("data", no_commit=True) + clean_staging() cache = tmp_dir / ".dvc" / "cache" - staging = tmp_dir / ".dvc" / "tmp" / _STAGING_DIR assert set(cache.glob("*/*")) == set() - assert set(staging.glob("*/*")) == { - staging / "1a" / "ca2c799df82929bbdd976557975546.dir" - } dvc.commit(os.path.join("data", "foo")) assert set(cache.glob("*/*")) == { cache / "1a" / "ca2c799df82929bbdd976557975546.dir", cache / "ac" / "bd18db4cc2f85cedef654fccc4a4d8", } - assert set(staging.glob("*/*")) == set() + clean_staging() dvc.commit(os.path.join("data", "subdir")) assert set(cache.glob("*/*")) == { @@ -158,7 +163,7 @@ def test_commit_granular_dir(tmp_dir, dvc): cache / "4c" / "e8d2a2cf314a52fa7f315ca37ca445", cache / "68" / "dde2c3c4e7953c2290f176bbdc9a54", } - assert set(staging.glob("*/*")) == set() + clean_staging() dvc.commit(os.path.join("data")) assert set(cache.glob("*/*")) == { @@ -168,13 +173,13 @@ def test_commit_granular_dir(tmp_dir, dvc): cache / "68" / "dde2c3c4e7953c2290f176bbdc9a54", cache / "37" / "b51d194a7513e45b56f6524f2d51f2", } - assert set(staging.glob("*/*")) == set() def test_commit_no_exec_missing_dep(tmp_dir, dvc): stage = dvc.run( name="my", cmd="mycmd", deps=["dep"], outs=["out"], no_exec=True ) + clean_staging() assert dvc.status(stage.path) with pytest.raises(DependencyDoesNotExistError): @@ -183,6 +188,7 @@ def test_commit_no_exec_missing_dep(tmp_dir, dvc): def test_commit_no_exec_missing_out(tmp_dir, dvc): stage = dvc.run(name="my", cmd="mycmd", outs=["out"], no_exec=True) + clean_staging() assert dvc.status(stage.path) with pytest.raises(OutputDoesNotExistError): @@ -192,6 +198,7 @@ def test_commit_no_exec_missing_out(tmp_dir, dvc): def test_commit_pipeline_stage(tmp_dir, dvc, run_copy): tmp_dir.gen("foo", "foo") stage = run_copy("foo", "bar", no_commit=True, name="copy-foo-bar") + clean_staging() assert dvc.status(stage.addressing) assert dvc.commit(stage.addressing, force=True) == [stage] assert not dvc.status(stage.addressing) @@ -205,6 +212,7 @@ def test_commit_pipeline_stage(tmp_dir, dvc, run_copy): def test_imported_entries_unchanged(tmp_dir, dvc, erepo_dir): with erepo_dir.chdir(): erepo_dir.dvc_gen("file", "file content", "initial commit") + clean_staging() stage = dvc.imp(os.fspath(erepo_dir), "file") diff --git a/tests/func/test_data_cloud.py b/tests/func/test_data_cloud.py index 39534deae..9878033b7 100644 --- a/tests/func/test_data_cloud.py +++ b/tests/func/test_data_cloud.py @@ -50,7 +50,7 @@ def test_cloud(tmp_dir, dvc, remote): # pylint:disable=unused-argument out = stage.outs[0] cache = out.cache_path foo_hash = out.hash_info - foo_objs = out.get_used_objs().get(None, set()) + foo_hashes = out.get_used_objs().get(None, set()) (stage_dir,) = tmp_dir.dvc_gen( { @@ -63,7 +63,6 @@ def test_cloud(tmp_dir, dvc, remote): # pylint:disable=unused-argument ) out_dir = stage_dir.outs[0] cache_dir = out_dir.cache_path - dir_objs = {out_dir.obj} dir_hash = out_dir.hash_info dir_hashes = {dir_hash} | { entry_obj.hash_info for _, entry_obj in out_dir.obj @@ -72,23 +71,23 @@ def test_cloud(tmp_dir, dvc, remote): # pylint:disable=unused-argument def _check_status(status, **kwargs): for key in ("ok", "missing", "new", "deleted"): expected = kwargs.get(key, set()) - assert expected == {obj.hash_info for obj in getattr(status, key)} + assert expected == set(getattr(status, key)) # Check status - status = dvc.cloud.status(foo_objs) + status = dvc.cloud.status(foo_hashes) _check_status(status, new={foo_hash}) - status_dir = dvc.cloud.status(dir_objs) + status_dir = dvc.cloud.status(dir_hashes) _check_status(status_dir, new=dir_hashes) # Move cache and check status # See issue https://github.com/iterative/dvc/issues/4383 for details backup_dir = dvc.odb.local.cache_dir + ".backup" move(dvc.odb.local.cache_dir, backup_dir) - status = dvc.cloud.status(foo_objs) + status = dvc.cloud.status(foo_hashes) _check_status(status, missing={foo_hash}) - status_dir = dvc.cloud.status(dir_objs) + status_dir = dvc.cloud.status(dir_hashes) _check_status(status_dir, missing=dir_hashes) # Restore original cache: @@ -96,42 +95,42 @@ def test_cloud(tmp_dir, dvc, remote): # pylint:disable=unused-argument move(backup_dir, dvc.odb.local.cache_dir) # Push and check status - dvc.cloud.push(foo_objs) + dvc.cloud.push(foo_hashes) assert os.path.exists(cache) assert os.path.isfile(cache) - dvc.cloud.push(dir_objs) + dvc.cloud.push(dir_hashes) assert os.path.isfile(cache_dir) - status = dvc.cloud.status(foo_objs) + status = dvc.cloud.status(foo_hashes) _check_status(status, ok={foo_hash}) - status_dir = dvc.cloud.status(dir_objs) + status_dir = dvc.cloud.status(dir_hashes) _check_status(status_dir, ok=dir_hashes) # Remove and check status remove(dvc.odb.local.cache_dir) - status = dvc.cloud.status(foo_objs) + status = dvc.cloud.status(foo_hashes) _check_status(status, deleted={foo_hash}) - status_dir = dvc.cloud.status(dir_objs) + status_dir = dvc.cloud.status(dir_hashes) _check_status(status_dir, deleted=dir_hashes) # Pull and check status - dvc.cloud.pull(foo_objs) + dvc.cloud.pull(foo_hashes) assert os.path.exists(cache) assert os.path.isfile(cache) with open(cache) as fd: assert fd.read() == "foo" - dvc.cloud.pull(dir_objs) + dvc.cloud.pull(dir_hashes) assert os.path.isfile(cache_dir) - status = dvc.cloud.status(foo_objs) + status = dvc.cloud.status(foo_hashes) _check_status(status, ok={foo_hash}) - status_dir = dvc.cloud.status(dir_objs) + status_dir = dvc.cloud.status(dir_hashes) _check_status(status_dir, ok=dir_hashes) @@ -242,10 +241,13 @@ def test_hash_recalculation(mocker, dvc, tmp_dir, local_remote): def test_missing_cache(tmp_dir, dvc, local_remote, caplog): + from tests.utils import clean_staging + tmp_dir.dvc_gen({"foo": "foo", "bar": "bar"}) # purge cache remove(dvc.odb.local.cache_dir) + clean_staging() header = ( "Some of the cache files do not exist " @@ -355,6 +357,8 @@ def test_pull_external_dvc_imports(tmp_dir, dvc, scm, erepo_dir): def clean(outs, dvc=None): + from tests.utils import clean_staging + if dvc: outs = outs + [dvc.odb.local.cache_dir] for path in outs: @@ -363,6 +367,7 @@ def clean(outs, dvc=None): if dvc: os.makedirs(dvc.odb.local.cache_dir, exist_ok=True) clean_repos() + clean_staging() def recurse_list_dir(d): diff --git a/tests/func/test_external_repo.py b/tests/func/test_external_repo.py index b74ff080c..1fcc9724d 100644 --- a/tests/func/test_external_repo.py +++ b/tests/func/test_external_repo.py @@ -3,13 +3,14 @@ from unittest.mock import ANY, patch from dvc.external_repo import CLONES, external_repo from dvc.fs.local import LocalFileSystem -from dvc.objects import save from dvc.objects.stage import stage +from dvc.objects.transfer import transfer from dvc.path_info import PathInfo from dvc.scm.git import Git from dvc.utils import relpath from dvc.utils.fs import makedirs, remove from tests.unit.fs.test_repo import make_subrepo +from tests.utils import clean_staging def test_external_repo(erepo_dir): @@ -198,16 +199,23 @@ def test_subrepos_are_ignored(tmp_dir, erepo_dir): # clear cache to test saving to cache cache_dir = tmp_dir / repo.odb.local.cache_dir remove(cache_dir) + clean_staging() makedirs(cache_dir) - obj = stage( + staging, obj = stage( repo.odb.local, PathInfo(repo.root_dir) / "dir", repo.repo_fs, "md5", dvcignore=repo.dvcignore, ) - save(repo.odb.local, obj) + transfer( + staging, + repo.odb.local, + {obj.hash_info}, + shallow=False, + move=True, + ) assert set(cache_dir.glob("??/*")) == { cache_dir / "e1" / "d9e8eae5374860ae025ec84cfd85c7.dir", cache_dir / "37" / "b51d194a7513e45b56f6524f2d51f2", diff --git a/tests/func/test_import.py b/tests/func/test_import.py index 87d7910e9..ccc99b4ca 100644 --- a/tests/func/test_import.py +++ b/tests/func/test_import.py @@ -13,6 +13,7 @@ from dvc.stage.exceptions import StagePathNotFoundError from dvc.system import System from dvc.utils.fs import makedirs, remove from tests.unit.fs.test_repo import make_subrepo +from tests.utils import clean_staging def test_import(tmp_dir, scm, dvc, erepo_dir): @@ -290,6 +291,8 @@ def test_push_wildcard_from_bare_git_repo( dvc_repo = make_tmp_dir("dvc-repo", scm=True, dvc=True) with dvc_repo.chdir(): dvc_repo.dvc.imp(os.fspath(tmp_dir), "dirextra") + clean_staging() + with pytest.raises(PathMissingError): dvc_repo.dvc.imp(os.fspath(tmp_dir), "dir123") @@ -593,6 +596,7 @@ def test_circular_import(tmp_dir, dvc, scm, erepo_dir): dvc.imp(os.fspath(erepo_dir), "dir", "dir_imported") scm.add("dir_imported.dvc") scm.commit("import") + clean_staging() with erepo_dir.chdir(): with pytest.raises(CircularImportError): diff --git a/tests/func/test_ls.py b/tests/func/test_ls.py index 911890cfd..1b716e957 100644 --- a/tests/func/test_ls.py +++ b/tests/func/test_ls.py @@ -512,6 +512,7 @@ def test_ls_target(erepo_dir, use_scm): ) def test_subrepo(dvc_top_level, erepo): from tests.func.test_get import make_subrepo + from tests.utils import clean_staging dvc_files = {"foo.txt": "foo.txt", "dvc_dir": {"lorem": "lorem"}} scm_files = {"bar.txt": "bar.txt", "scm_dir": {"ipsum": "ipsum"}} @@ -523,6 +524,7 @@ def test_subrepo(dvc_top_level, erepo): repo.scm_gen(scm_files, commit=f"scm track for top {repo}") if hasattr(repo, "dvc"): repo.dvc_gen(dvc_files, commit=f"dvc track for {repo}") + clean_staging() def _list_files(repo, path=None): return set(map(itemgetter("path"), Repo.ls(os.fspath(repo), path))) diff --git a/tests/func/test_remote.py b/tests/func/test_remote.py index 5daafa569..edfc66d85 100644 --- a/tests/func/test_remote.py +++ b/tests/func/test_remote.py @@ -157,18 +157,16 @@ def test_dir_hash_should_be_key_order_agnostic(tmp_dir, dvc): ) tree.digest() with patch("dvc.objects.stage._get_tree_obj", return_value=tree): - hash1 = stage( - dvc.odb.local, path_info, dvc.odb.local.fs, "md5" - ).hash_info + _, obj = stage(dvc.odb.local, path_info, dvc.odb.local.fs, "md5") + hash1 = obj.hash_info tree = Tree.from_list( [{"md5": "1", "relpath": "1"}, {"md5": "2", "relpath": "2"}] ) tree.digest() with patch("dvc.objects.stage._get_tree_obj", return_value=tree): - hash2 = stage( - dvc.odb.local, path_info, dvc.odb.local.fs, "md5" - ).hash_info + _, obj = stage(dvc.odb.local, path_info, dvc.odb.local.fs, "md5") + hash2 = obj.hash_info assert hash1 == hash2 diff --git a/tests/func/test_run_single_stage.py b/tests/func/test_run_single_stage.py index 146d0a020..a443dc349 100644 --- a/tests/func/test_run_single_stage.py +++ b/tests/func/test_run_single_stage.py @@ -900,10 +900,13 @@ class TestShouldNotCheckoutUponCorruptedLocalHardlinkCache(TestDvc): self.dvc = DvcRepo(".") def test(self): + from tests.utils import clean_staging + cmd = f"python {self.CODE} {self.FOO} {self.BAR}" stage = self.dvc.run( deps=[self.FOO], outs=[self.BAR], cmd=cmd, single_stage=True ) + clean_staging() os.chmod(self.BAR, 0o644) with open(self.BAR, "w") as fd: diff --git a/tests/unit/fs/test_dvc.py b/tests/unit/fs/test_dvc.py index 5c2a4f70c..9c0b789e3 100644 --- a/tests/unit/fs/test_dvc.py +++ b/tests/unit/fs/test_dvc.py @@ -243,15 +243,15 @@ def test_get_hash_granular(tmp_dir, dvc): fs = DvcFileSystem(repo=dvc) subdir = PathInfo(tmp_dir) / "dir" / "subdir" assert fs.info(subdir).get("md5") is None - assert stage(dvc.odb.local, subdir, fs, "md5").hash_info == HashInfo( + _, obj = stage(dvc.odb.local, subdir, fs, "md5", dry_run=True) + assert obj.hash_info == HashInfo( "md5", "af314506f1622d107e0ed3f14ec1a3b5.dir" ) assert ( fs.info(subdir / "data")["md5"] == "8d777f385d3dfec8815d20f7496026dc" ) - assert stage( - dvc.odb.local, subdir / "data", fs, "md5" - ).hash_info == HashInfo("md5", "8d777f385d3dfec8815d20f7496026dc") + _, obj = stage(dvc.odb.local, subdir / "data", fs, "md5", dry_run=True) + assert obj.hash_info == HashInfo("md5", "8d777f385d3dfec8815d20f7496026dc") def test_get_hash_dirty_file(tmp_dir, dvc): @@ -261,9 +261,10 @@ def test_get_hash_dirty_file(tmp_dir, dvc): fs = DvcFileSystem(repo=dvc) expected = "8c7dd922ad47494fc02c388e12c00eac" assert fs.info(PathInfo(tmp_dir) / "file").get("md5") == expected - assert stage( - dvc.odb.local, PathInfo(tmp_dir) / "file", fs, "md5" - ).hash_info == HashInfo("md5", expected) + _, obj = stage( + dvc.odb.local, PathInfo(tmp_dir) / "file", fs, "md5", dry_run=True + ) + assert obj.hash_info == HashInfo("md5", expected) def test_get_hash_dirty_dir(tmp_dir, dvc): @@ -273,6 +274,7 @@ def test_get_hash_dirty_dir(tmp_dir, dvc): fs = DvcFileSystem(repo=dvc) expected = "5ea40360f5b4ec688df672a4db9c17d1.dir" assert fs.info(PathInfo(tmp_dir) / "dir").get("md5") == expected - assert stage( - dvc.odb.local, PathInfo(tmp_dir) / "dir", fs, "md5" - ).hash_info == HashInfo("md5", expected) + _, obj = stage( + dvc.odb.local, PathInfo(tmp_dir) / "dir", fs, "md5", dry_run=True + ) + assert obj.hash_info == HashInfo("md5", expected) diff --git a/tests/unit/fs/test_repo.py b/tests/unit/fs/test_repo.py index bab74f2ed..6115797f3 100644 --- a/tests/unit/fs/test_repo.py +++ b/tests/unit/fs/test_repo.py @@ -8,6 +8,7 @@ from dvc.fs.repo import RepoFileSystem from dvc.hash_info import HashInfo from dvc.objects.stage import stage from dvc.path_info import PathInfo +from tests.utils import clean_staging def test_exists(tmp_dir, dvc): @@ -512,9 +513,8 @@ def test_get_hash_cached_file(tmp_dir, dvc, mocker): fs = RepoFileSystem(repo=dvc) expected = "acbd18db4cc2f85cedef654fccc4a4d8" assert fs.info(PathInfo(tmp_dir) / "foo").get("md5") is None - assert stage( - dvc.odb.local, PathInfo(tmp_dir) / "foo", fs, "md5" - ).hash_info == HashInfo("md5", expected) + _, obj = stage(dvc.odb.local, PathInfo(tmp_dir) / "foo", fs, "md5") + assert obj.hash_info == HashInfo("md5", expected) (tmp_dir / "foo").unlink() assert fs.info(PathInfo(tmp_dir) / "foo")["md5"] == expected @@ -526,15 +526,17 @@ def test_get_hash_cached_dir(tmp_dir, dvc, mocker): fs = RepoFileSystem(repo=dvc) expected = "8761c4e9acad696bee718615e23e22db.dir" assert fs.info(PathInfo(tmp_dir) / "dir").get("md5") is None - assert stage( - dvc.odb.local, PathInfo(tmp_dir) / "dir", fs, "md5" - ).hash_info == HashInfo("md5", "8761c4e9acad696bee718615e23e22db.dir") + _, obj = stage(dvc.odb.local, PathInfo(tmp_dir) / "dir", fs, "md5") + assert obj.hash_info == HashInfo( + "md5", "8761c4e9acad696bee718615e23e22db.dir" + ) shutil.rmtree(tmp_dir / "dir") assert fs.info(PathInfo(tmp_dir) / "dir")["md5"] == expected - assert stage( - dvc.odb.local, PathInfo(tmp_dir) / "dir", fs, "md5" - ).hash_info == HashInfo("md5", "8761c4e9acad696bee718615e23e22db.dir") + _, obj = stage(dvc.odb.local, PathInfo(tmp_dir) / "dir", fs, "md5") + assert obj.hash_info == HashInfo( + "md5", "8761c4e9acad696bee718615e23e22db.dir" + ) def test_get_hash_cached_granular(tmp_dir, dvc, mocker): @@ -544,13 +546,13 @@ def test_get_hash_cached_granular(tmp_dir, dvc, mocker): fs = RepoFileSystem(repo=dvc) subdir = PathInfo(tmp_dir) / "dir" / "subdir" assert fs.info(subdir).get("md5") is None - assert stage(dvc.odb.local, subdir, fs, "md5").hash_info == HashInfo( + _, obj = stage(dvc.odb.local, subdir, fs, "md5") + assert obj.hash_info == HashInfo( "md5", "af314506f1622d107e0ed3f14ec1a3b5.dir" ) assert fs.info(subdir / "data").get("md5") is None - assert stage( - dvc.odb.local, subdir / "data", fs, "md5" - ).hash_info == HashInfo("md5", "8d777f385d3dfec8815d20f7496026dc") + _, obj = stage(dvc.odb.local, subdir / "data", fs, "md5") + assert obj.hash_info == HashInfo("md5", "8d777f385d3dfec8815d20f7496026dc") (tmp_dir / "dir" / "subdir" / "data").unlink() assert ( fs.info(subdir / "data")["md5"] == "8d777f385d3dfec8815d20f7496026dc" @@ -568,50 +570,66 @@ def test_get_hash_mixed_dir(tmp_dir, scm, dvc): ] ) tmp_dir.scm.commit("add dir") + clean_staging() fs = RepoFileSystem(repo=dvc) - actual = stage( - dvc.odb.local, PathInfo(tmp_dir) / "dir", fs, "md5" - ).hash_info - expected = HashInfo("md5", "e1d9e8eae5374860ae025ec84cfd85c7.dir") - assert actual == expected + _, obj = stage(dvc.odb.local, PathInfo(tmp_dir) / "dir", fs, "md5") + assert obj.hash_info == HashInfo( + "md5", "e1d9e8eae5374860ae025ec84cfd85c7.dir" + ) def test_get_hash_dirty_file(tmp_dir, dvc): - tmp_dir.dvc_gen("file", "file") - (tmp_dir / "file").write_text("something") + from dvc.objects import check + from dvc.objects.errors import ObjectFormatError + from dvc.objects.stage import get_file_hash + tmp_dir.dvc_gen("file", "file") + file_hash_info = HashInfo("md5", "8c7dd922ad47494fc02c388e12c00eac") + + (tmp_dir / "file").write_text("something") + something_hash_info = HashInfo("md5", "437b930db84b8079c2dd804a71936b5f") + + clean_staging() + + # file is modified in workspace + # get_file_hash(file) should return workspace hash, not DVC cached hash fs = RepoFileSystem(repo=dvc) assert fs.info(PathInfo(tmp_dir) / "file").get("md5") is None - actual = stage( - dvc.odb.local, PathInfo(tmp_dir) / "file", fs, "md5" - ).hash_info - expected = HashInfo("md5", "437b930db84b8079c2dd804a71936b5f") - assert actual == expected + staging, obj = stage(dvc.odb.local, PathInfo(tmp_dir) / "file", fs, "md5") + assert obj.hash_info == something_hash_info + check(staging, obj) + # file is removed in workspace + # any staged object referring to modified workspace obj is now invalid (tmp_dir / "file").unlink() - assert ( - fs.info(PathInfo(tmp_dir) / "file")["md5"] - == "8c7dd922ad47494fc02c388e12c00eac" + with pytest.raises(ObjectFormatError): + check(staging, obj) + + # get_file_hash(file) should return DVC cached hash + assert fs.info(PathInfo(tmp_dir) / "file")["md5"] == file_hash_info.value + hash_info = get_file_hash( + PathInfo(tmp_dir) / "file", fs, "md5", state=dvc.state ) - actual = stage( - dvc.odb.local, PathInfo(tmp_dir) / "file", fs, "md5" - ).hash_info - expected = HashInfo("md5", "8c7dd922ad47494fc02c388e12c00eac") - assert actual == expected + assert hash_info == file_hash_info + + # tmp_dir/file can be staged even though it is missing in workspace since + # repofs will use the DVC cached hash (and refer to the local cache object) + _, obj = stage(dvc.odb.local, PathInfo(tmp_dir) / "file", fs, "md5") + assert obj.hash_info == file_hash_info def test_get_hash_dirty_dir(tmp_dir, dvc): tmp_dir.dvc_gen({"dir": {"foo": "foo", "bar": "bar"}}) (tmp_dir / "dir" / "baz").write_text("baz") + clean_staging() fs = RepoFileSystem(repo=dvc) - actual = stage( - dvc.odb.local, PathInfo(tmp_dir) / "dir", fs, "md5" - ).hash_info - expected = HashInfo("md5", "ba75a2162ca9c29acecb7957105a0bc2.dir") - assert actual == expected - assert actual.nfiles == 3 + _, obj = stage(dvc.odb.local, PathInfo(tmp_dir) / "dir", fs, "md5") + assert obj.hash_info == HashInfo( + "md5", "ba75a2162ca9c29acecb7957105a0bc2.dir" + ) + assert obj.hash_info.nfiles == 3 @pytest.mark.parametrize("traverse_subrepos", [True, False]) diff --git a/tests/unit/objects/db/test_local.py b/tests/unit/objects/db/test_local.py index 98d23ea71..da305b9a1 100644 --- a/tests/unit/objects/db/test_local.py +++ b/tests/unit/objects/db/test_local.py @@ -6,7 +6,6 @@ import pytest from dvc.fs.local import LocalFileSystem from dvc.hash_info import HashInfo from dvc.objects.db.local import LocalObjectDB -from dvc.objects.file import HashFile from dvc.path_info import PathInfo @@ -18,21 +17,17 @@ def test_status_download_optimization(mocker, dvc): from dvc.objects.status import compare_status odb = LocalObjectDB(LocalFileSystem(), PathInfo(".")) - objs = { - HashFile( - None, odb.fs, HashInfo("md5", "acbd18db4cc2f85cedef654fccc4a4d8") - ), - HashFile( - None, odb.fs, HashInfo("md5", "37b51d194a7513e45b56f6524f2d51f2") - ), + obj_ids = { + HashInfo("md5", "acbd18db4cc2f85cedef654fccc4a4d8"), + HashInfo("md5", "37b51d194a7513e45b56f6524f2d51f2"), } - local_exists = [obj.hash_info.value for obj in objs] + local_exists = [hash_info.value for hash_info in obj_ids] mocker.patch.object(odb, "hashes_exist", return_value=local_exists) src_odb = mocker.Mock() - compare_status(src_odb, odb, objs, check_deleted=False) + compare_status(src_odb, odb, obj_ids, check_deleted=False) assert src_odb.hashes_exist.call_count == 0 @@ -91,28 +86,24 @@ def test_set_exec_ignore_errors(tmp_dir, dvc, mocker, err): def test_staging_file(tmp_dir, dvc): - from dvc.objects import check, save - from dvc.objects.stage import get_staging, stage + from dvc.objects import check + from dvc.objects.stage import stage + from dvc.objects.transfer import transfer tmp_dir.gen("foo", "foo") fs = LocalFileSystem() local_odb = dvc.odb.local - staging_odb = get_staging(local_odb) - obj = stage(local_odb, tmp_dir / "foo", fs, "md5") + staging_odb, obj = stage(local_odb, tmp_dir / "foo", fs, "md5") - for odb in (local_odb, staging_odb): - path_info = odb.hash_to_path_info(obj.hash_info.value) - assert not odb.fs.exists(path_info) + assert not local_odb.exists(obj.hash_info) + assert staging_odb.exists(obj.hash_info) - # check for file after staging should fail since files are not added on - # stage() with pytest.raises(FileNotFoundError): check(local_odb, obj) - with pytest.raises(FileNotFoundError): - check(staging_odb, obj) + check(staging_odb, obj) - save(local_odb, obj) + transfer(staging_odb, local_odb, {obj.hash_info}, move=True) check(local_odb, obj) with pytest.raises(FileNotFoundError): check(staging_odb, obj) @@ -122,35 +113,27 @@ def test_staging_file(tmp_dir, dvc): def test_staging_dir(tmp_dir, dvc): - from dvc.objects import check, save - from dvc.objects.stage import get_staging, stage + from dvc.objects import check + from dvc.objects.stage import stage + from dvc.objects.transfer import transfer tmp_dir.gen({"dir": {"foo": "foo", "bar": "bar"}}) fs = LocalFileSystem() local_odb = dvc.odb.local - staging_odb = get_staging(local_odb) - obj = stage(local_odb, tmp_dir / "dir", fs, "md5") + staging_odb, obj = stage(local_odb, tmp_dir / "dir", fs, "md5") - path_info = local_odb.hash_to_path_info(obj.hash_info.value) - assert not local_odb.fs.exists(path_info) - path_info = staging_odb.hash_to_path_info(obj.hash_info.value) - assert staging_odb.fs.exists(path_info) + assert not local_odb.exists(obj.hash_info) + assert staging_odb.exists(obj.hash_info) - # check for raw object after staging should pass only when using the - # staging odb - raw = HashFile(obj.path_info, obj.fs, obj.hash_info) with pytest.raises(FileNotFoundError): - check(local_odb, raw) - check(staging_odb, raw) + check(local_odb, obj) + check(staging_odb, obj) - # checking the entire tree should fail since individual file entries are - # not added on stage() + transfer(staging_odb, local_odb, {obj.hash_info}, shallow=False, move=True) + check(local_odb, obj) with pytest.raises(FileNotFoundError): check(staging_odb, obj) - save(local_odb, obj) - check(local_odb, obj) - path_info = local_odb.hash_to_path_info(obj.hash_info.value) assert fs.exists(path_info) diff --git a/tests/unit/repo/test_repo.py b/tests/unit/repo/test_repo.py index 316a5411a..56e07c8ba 100644 --- a/tests/unit/repo/test_repo.py +++ b/tests/unit/repo/test_repo.py @@ -45,8 +45,6 @@ def test_find_outs_by_path_does_graph_checks(tmp_dir, dvc): [os.path.join("dir", "subdir", "file"), os.path.join("dir", "subdir")], ) def test_used_objs(tmp_dir, dvc, path): - from dvc.objects.tree import Tree - tmp_dir.dvc_gen({"dir": {"subdir": {"file": "file"}, "other": "other"}}) expected = { @@ -55,11 +53,8 @@ def test_used_objs(tmp_dir, dvc, path): } used = set() - for _, objs in dvc.used_objs([path]).items(): - for obj in objs: - used.add(obj.hash_info) - if isinstance(obj, Tree): - used.update(entry_obj.hash_info for _, entry_obj in obj) + for _, obj_ids in dvc.used_objs([path]).items(): + used.update(obj_ids) assert used == expected diff --git a/tests/unit/utils/test_fs.py b/tests/unit/utils/test_fs.py index 23eb644d5..d56e2e1b0 100644 --- a/tests/unit/utils/test_fs.py +++ b/tests/unit/utils/test_fs.py @@ -38,11 +38,11 @@ class TestMtimeAndSize(TestDir): ) self.assertIs(type(file_time), str) - self.assertIs(type(file_size), str) - self.assertEqual(file_size, str(actual_file_size)) + self.assertIs(type(file_size), int) + self.assertEqual(file_size, actual_file_size) self.assertIs(type(dir_time), str) - self.assertIs(type(dir_size), str) - self.assertEqual(dir_size, str(actual_dir_size)) + self.assertIs(type(dir_size), int) + self.assertEqual(dir_size, actual_dir_size) class TestContainsLink(TestCase): diff --git a/tests/utils/__init__.py b/tests/utils/__init__.py index dbd35a84b..c14cc7ae3 100644 --- a/tests/utils/__init__.py +++ b/tests/utils/__init__.py @@ -42,3 +42,15 @@ def dump_sv(stream, metrics, delimiter=",", header=True): writer = csv.writer(stream) for d in metrics: writer.writerow(list(d.values())) + + +def clean_staging(): + from dvc.fs.memory import MemoryFileSystem + from dvc.objects.stage import _STAGING_MEMFS_PATH + + try: + MemoryFileSystem().fs.rm( + f"memory://{_STAGING_MEMFS_PATH}", recursive=True + ) + except FileNotFoundError: + pass