import/get: support pre-fetching LFS objects from Git-LFS repos

This commit is contained in:
Peter Rowlands
2023-11-29 14:39:54 +09:00
committed by Peter Rowlands (변기호)
parent 7edcb6b1e5
commit b0d0315689
3 changed files with 54 additions and 2 deletions
+7 -1
View File
@@ -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(
+29
View File
@@ -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,
)
+18 -1
View File
@@ -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()