diff --git a/dvc/fs/__init__.py b/dvc/fs/__init__.py index 77ac37cab..1e8627efc 100644 --- a/dvc/fs/__init__.py +++ b/dvc/fs/__init__.py @@ -5,6 +5,8 @@ from dvc_http import HTTPFileSystem, HTTPSFileSystem # noqa: F401 from dvc.config import ConfigError as RepoConfigError from dvc.config_schema import SCHEMA, Invalid + +# pylint: disable=unused-import from dvc_objects.fs import ( # noqa: F401 LocalFileSystem, MemoryFileSystem, @@ -27,7 +29,7 @@ from dvc_objects.fs.path import Path # noqa: F401 from .callbacks import Callback from .data import DataFileSystem # noqa: F401 -from .dvc import DVCFileSystem # noqa: F401 +from .dvc import DVCFileSystem from .git import GitFileSystem # noqa: F401 known_implementations.update( @@ -47,6 +49,8 @@ known_implementations.update( def download( fs: "FileSystem", fs_path: str, to: str, jobs: Optional[int] = None ) -> int: + from dvc.scm import lfs_prefetch + with Callback.as_tqdm_callback( desc=f"Downloading {fs.path.name(fs_path)}", unit="files", @@ -70,6 +74,8 @@ def download( from_infos = [fs_path] to_infos = [to] + if isinstance(fs, DVCFileSystem): + lfs_prefetch(fs, from_infos) cb.set_size(len(from_infos)) jobs = jobs or fs.jobs generic.copy( diff --git a/dvc/scm.py b/dvc/scm.py index 323adaf59..435b34a51 100644 --- a/dvc/scm.py +++ b/dvc/scm.py @@ -24,6 +24,8 @@ from dvc.progress import Tqdm if TYPE_CHECKING: from scmrepo.progress import GitProgressEvent + from dvc.fs import FileSystem + class SCMError(DvcException): """Base class for source control management errors.""" @@ -262,3 +264,30 @@ def iter_revs( rev_resolver = partial(resolve_rev, scm) return group_by(rev_resolver, results) + + +def lfs_prefetch(fs: "FileSystem", paths: List[str]): + from scmrepo.git.lfs import fetch as _lfs_fetch + + from dvc.fs.dvc import DVCFileSystem + from dvc.fs.git import GitFileSystem + + if isinstance(fs, DVCFileSystem) and isinstance(fs.repo.fs, GitFileSystem): + git_fs = fs.repo.fs + scm = fs.repo.scm + assert isinstance(scm, Git) + else: + return + + try: + if "filter=lfs" not in git_fs.open(".gitattributes").read(): + return + except OSError: + return + with TqdmGit(desc="Checking for Git-LFS objects") as pbar: + _lfs_fetch( + scm, + [git_fs.rev], + include=[(path if path.startswith("/") else f"/{path}") for path in paths], + progress=pbar.update_git, + ) diff --git a/tests/func/test_scm.py b/tests/func/test_scm.py index f7a227557..b5cab8aea 100644 --- a/tests/func/test_scm.py +++ b/tests/func/test_scm.py @@ -3,7 +3,7 @@ import os import pytest from git import Repo -from dvc.scm import SCM, Git, NoSCM, SCMError +from dvc.scm import SCM, Git, NoSCM, SCMError, lfs_prefetch def test_init_none(tmp_dir): @@ -27,3 +27,20 @@ def test_init_sub_dir(tmp_dir): scm = SCM(os.fspath(subdir)) assert scm.root_dir == os.fspath(tmp_dir) + + +def test_lfs_prefetch(tmp_dir, dvc, scm, mocker): + mock_fetch = mocker.patch("scmrepo.git.lfs.fetch") + rev = scm.get_rev() + + with dvc.switch(rev): + lfs_prefetch(dvc.dvcfs, ["foo"]) + mock_fetch.assert_not_called() + + tmp_dir.scm_gen( + ".gitattributes", ".lfs filter=lfs diff=lfs merge=lfs -text", commit="init lfs" + ) + rev = scm.get_rev() + with dvc.switch(rev): + lfs_prefetch(dvc.dvcfs, ["foo"]) + mock_fetch.assert_called_once()