scmrepo 1.4.0__py3-none-any.whl → 1.5.0__py3-none-any.whl

scmrepo/git/lfs/client.py

@@ -0,0 +1,223 @@
+import logging
+from contextlib import AbstractContextManager
+from functools import wraps
+from typing import TYPE_CHECKING, Any, Awaitable, Callable, Dict, Iterable, Optional
+
+import aiohttp
+from dvc_http import HTTPFileSystem
+from dvc_http.retry import ReadOnlyRetryClient
+from dvc_objects.executors import batch_coros
+from dvc_objects.fs import localfs
+from dvc_objects.fs.callbacks import DEFAULT_CALLBACK
+from dvc_objects.fs.utils import as_atomic
+from fsspec.asyn import sync_wrapper
+from funcy import cached_property
+
+from ..credentials import Credential, CredentialNotFoundError
+from .exceptions import LFSError
+from .pointer import Pointer
+
+if TYPE_CHECKING:
+    from dvc_objects.fs.callbacks import Callback
+
+    from .storage import LFSStorage
+
+logger = logging.getLogger(__name__)
+
+
+class _LFSClient(ReadOnlyRetryClient):
+    async def _request(self, *args, **kwargs):
+        return await super()._request(*args, **kwargs)  # pylint: disable=no-member
+
+
+# pylint: disable=abstract-method
+class _LFSFileSystem(HTTPFileSystem):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def _prepare_credentials(self, **config):
+        return {}
+
+    async def get_client(self, **kwargs):
+        from aiohttp_retry import ExponentialRetry
+        from dvc_http import make_context
+
+        kwargs["retry_options"] = ExponentialRetry(
+            attempts=self.SESSION_RETRIES,
+            factor=self.SESSION_BACKOFF_FACTOR,
+            max_timeout=self.REQUEST_TIMEOUT,
+            exceptions={aiohttp.ClientError},
+        )
+
+        # The default total timeout for an aiohttp request is 300 seconds
+        # which is too low for DVC's interactions when dealing with large
+        # data blobs. We remove the total timeout, and only limit the time
+        # that is spent when connecting to the remote server and waiting
+        # for new data portions.
+        connect_timeout = kwargs.pop("connect_timeout")
+        kwargs["timeout"] = aiohttp.ClientTimeout(
+            total=None,
+            connect=connect_timeout,
+            sock_connect=connect_timeout,
+            sock_read=kwargs.pop("read_timeout"),
+        )
+
+        kwargs["connector"] = aiohttp.TCPConnector(
+            # Force cleanup of closed SSL transports.
+            # See https://github.com/iterative/dvc/issues/7414
+            enable_cleanup_closed=True,
+            ssl=make_context(kwargs.pop("ssl_verify", None)),
+        )
+
+        return ReadOnlyRetryClient(**kwargs)
+
+
+def _authed(f: Callable[..., Awaitable]):
+    """Set credentials and retry the given coroutine if needed."""
+
+    # pylint: disable=protected-access
+    @wraps(f)  # type: ignore[arg-type]
+    async def wrapper(self, *args, **kwargs):
+        try:
+            return await f(self, *args, **kwargs)
+        except aiohttp.ClientResponseError as exc:
+            if exc.status != 401:
+                raise
+            session = await self._set_session()
+            if session.auth:
+                raise
+            auth = self._get_auth()
+            if auth is None:
+                raise
+            self._session._auth = auth
+        return await f(self, *args, **kwargs)
+
+    return wrapper
+
+
+class LFSClient(AbstractContextManager):
+    """Naive read-only LFS HTTP client."""
+
+    JSON_CONTENT_TYPE = "application/vnd.git-lfs+json"
+
+    def __init__(
+        self,
+        url: str,
+        git_url: Optional[str] = None,
+        headers: Optional[Dict[str, str]] = None,
+    ):
+        """
+        Args:
+            url: LFS server URL.
+        """
+        self.url = url
+        self.git_url = git_url
+        self.headers: Dict[str, str] = headers or {}
+
+    def __exit__(self, *args, **kwargs):
+        self.close()
+
+    @cached_property
+    def fs(self) -> "_LFSFileSystem":
+        return _LFSFileSystem()
+
+    @property
+    def httpfs(self) -> "HTTPFileSystem":
+        return self.fs.fs
+
+    @property
+    def loop(self):
+        return self.httpfs.loop
+
+    @classmethod
+    def from_git_url(cls, git_url: str) -> "LFSClient":
+        if git_url.endswith(".git"):
+            url = f"{git_url}/info/lfs"
+        else:
+            url = f"{git_url}.git/info/lfs"
+        return cls(url, git_url=git_url)
+
+    def close(self):
+        pass
+
+    def _get_auth(self) -> Optional[aiohttp.BasicAuth]:
+        try:
+            creds = Credential(url=self.git_url).fill()
+            if creds.username and creds.password:
+                return aiohttp.BasicAuth(creds.username, creds.password)
+        except CredentialNotFoundError:
+            pass
+        return None
+
+    async def _set_session(self) -> aiohttp.ClientSession:
+        return await self.fs.fs.set_session()
+
+    @_authed
+    async def _batch_request(
+        self,
+        objects: Iterable[Pointer],
+        upload: bool = False,
+        ref: Optional[str] = None,
+        hash_algo: str = "sha256",
+    ) -> Dict[str, Any]:
+        """Send LFS API /objects/batch request."""
+        url = f"{self.url}/objects/batch"
+        body: Dict[str, Any] = {
+            "operation": "upload" if upload else "download",
+            "transfers": ["basic"],
+            "objects": [{"oid": obj.oid, "size": obj.size} for obj in objects],
+            "hash_algo": hash_algo,
+        }
+        if ref:
+            body["ref"] = [{"name": ref}]
+        session = await self._set_session()
+        headers = dict(self.headers)
+        headers["Content-Type"] = self.JSON_CONTENT_TYPE
+        async with session.post(
+            url,
+            headers=headers,
+            json=body,
+        ) as resp:
+            data = await resp.json()
+        return data
+
+    @_authed
+    async def _download(
+        self,
+        storage: "LFSStorage",
+        objects: Iterable[Pointer],
+        callback: "Callback" = DEFAULT_CALLBACK,
+        **kwargs,
+    ):
+        async def _get_one(from_path: str, to_path: str, **kwargs):
+            get_coro = callback.wrap_and_branch_coro(
+                self.httpfs._get_file  # pylint: disable=protected-access
+            )
+            with as_atomic(localfs, to_path, create_parents=True) as tmp_file:
+                return await get_coro(
+                    from_path,
+                    tmp_file,
+                    **kwargs,
+                )
+
+        resp_data = await self._batch_request(objects, **kwargs)
+        if resp_data.get("transfer") != "basic":
+            raise LFSError("Unsupported LFS transfer type")
+        coros = []
+        for data in resp_data.get("objects", []):
+            obj = Pointer(data["oid"], data["size"])
+            download = data.get("actions", {}).get("download", {})
+            url = download.get("href")
+            if not url:
+                logger.debug("No download URL for LFS object '%s'", obj)
+                continue
+            headers = download.get("header", {})
+            to_path = storage.oid_to_path(obj.oid)
+            coros.append(_get_one(url, to_path, headers=headers))
+        for result in await batch_coros(
+            coros, batch_size=self.fs.jobs, return_exceptions=True
+        ):
+            if isinstance(result, BaseException):
+                raise result
+
+    download = sync_wrapper(_download)

scmrepo/git/lfs/exceptions.py

@@ -0,0 +1,5 @@
+from scmrepo.exceptions import SCMError
+
+
+class LFSError(SCMError):
+    pass

scmrepo/git/lfs/fetch.py

@@ -0,0 +1,162 @@
+import fnmatch
+import io
+import os
+from typing import TYPE_CHECKING, Callable, Iterable, Iterator, List, Optional, Set
+
+from scmrepo.exceptions import InvalidRemote, SCMError
+
+from .pointer import HEADERS, Pointer
+
+if TYPE_CHECKING:
+    from scmrepo.git import Git
+    from scmrepo.git.config import Config
+    from scmrepo.progress import GitProgressEvent
+
+
+def fetch(
+    scm: "Git",
+    revs: Optional[List[str]] = None,
+    remote: Optional[str] = None,
+    include: Optional[List[str]] = None,
+    exclude: Optional[List[str]] = None,
+    progress: Optional[Callable[["GitProgressEvent"], None]] = None,
+):
+    # NOTE: This currently does not support fetching objects from the worktree
+    if not revs:
+        revs = ["HEAD"]
+    objects: Set[Pointer] = set()
+    for rev in revs:
+        objects.update(
+            pointer
+            for pointer in _collect_objects(scm, rev, include, exclude)
+            if not scm.lfs_storage.exists(pointer)
+        )
+    if not objects:
+        return
+    try:
+        url = get_fetch_url(scm, remote=remote)
+    except InvalidRemote:
+        if remote:
+            # treat remote as a raw Git remote
+            url = remote
+        else:
+            raise
+    scm.lfs_storage.fetch(url, objects, progress=progress)
+
+
+def get_fetch_url(scm: "Git", remote: Optional[str] = None):  # noqa: C901
+    """Return LFS fetch URL for the specified repository."""
+    git_config = scm.get_config()
+
+    # check lfs.url (can be set in git config and .lfsconfig)
+    try:
+        return git_config.get(("lfs",), "url")
+    except KeyError:
+        pass
+    try:
+        lfs_config: Optional["Config"] = scm.get_config(
+            os.path.join(scm.root_dir, ".lfsconfig")
+        )
+    except FileNotFoundError:
+        lfs_config = None
+    if lfs_config:
+        try:
+            return lfs_config.get(("lfs",), "url")
+        except KeyError:
+            pass
+
+    # use:
+    #   current tracking-branch remote
+    #   or remote.lfsdefault  (can only be set in git config)
+    #   or "origin"
+    # in that order
+    if not remote:
+        try:
+            remote = scm.active_branch_remote()
+        except SCMError:
+            pass
+    if not remote:
+        try:
+            remote = git_config.get(("remote",), "lfsdefault")
+        except KeyError:
+            remote = "origin"
+
+    # check remote.*.lfsurl (can be set in git config and .lfsconfig)
+    assert remote is not None
+    try:
+        return git_config.get(("remote", remote), "lfsurl")
+    except KeyError:
+        pass
+    if lfs_config:
+        try:
+            return lfs_config.get(("remote", remote), "lfsurl")
+        except KeyError:
+            pass
+
+    # return default Git fetch URL for this remote
+    return scm.get_remote_url(remote)
+
+
+def _collect_objects(
+    scm: "Git",
+    rev: str,
+    include: Optional[List[str]],
+    exclude: Optional[List[str]],
+) -> Iterator[Pointer]:
+    fs = scm.get_fs(rev)
+    for path in _filter_paths(fs.find("/"), include, exclude):
+        check_path = path.lstrip("/")
+        if scm.check_attr(check_path, "filter", source=rev) == "lfs":
+            try:
+                with fs.open(path, "rb", raw=True) as fobj:
+                    with io.BufferedReader(fobj) as reader:
+                        data = reader.peek(100)
+                        if any(data.startswith(header) for header in HEADERS):
+                            yield Pointer.load(reader)
+            except (ValueError, OSError):
+                pass
+
+
+def _filter_paths(
+    paths: Iterable[str], include: Optional[List[str]], exclude: Optional[List[str]]
+) -> Iterator[str]:
+    filtered = set()
+    if include:
+        for pattern in include:
+            filtered.update(fnmatch.filter(paths, pattern))
+    else:
+        filtered.update(paths)
+    if exclude:
+        for pattern in exclude:
+            filtered.difference_update(fnmatch.filter(paths, pattern))
+    yield from filtered
+
+
+if __name__ == "__main__":
+    # Minimal `git lfs fetch` CLI implementation
+    import argparse
+    import sys
+
+    from scmrepo.git import Git  # noqa: F811
+
+    parser = argparse.ArgumentParser(
+        description=(
+            "Download Git LFS objects at the given refs from the specified remote."
+        ),
+    )
+    parser.add_argument(
+        "remote",
+        nargs="?",
+        default="origin",
+        help="Remote to fetch from. Defaults to 'origin'.",
+    )
+    parser.add_argument(
+        "refs",
+        nargs="*",
+        default=["HEAD"],
+        help="Refs or commits to fetch. Defaults to 'HEAD'.",
+    )
+    args = parser.parse_args()
+    with Git(".") as scm_:  # pylint: disable=E0601
+        print("fetch: fetching reference", ", ".join(args.refs), file=sys.stderr)
+        fetch(scm_, revs=args.refs, remote=args.remote)

scmrepo/git/lfs/object.py

@@ -0,0 +1,15 @@
+from dataclasses import dataclass, fields
+from typing import Any, Dict
+
+
+@dataclass(frozen=True)
+class LFSObject:
+    oid: str
+    size: int
+
+    def __str__(self) -> str:
+        return self.oid
+
+    @classmethod
+    def from_dict(cls, d: Dict[str, Any]) -> "LFSObject":
+        return cls(**{k: v for k, v in d.items() if k in fields(cls)})

scmrepo/git/lfs/pointer.py

@@ -0,0 +1,109 @@
+import hashlib
+import io
+import logging
+from dataclasses import dataclass
+from typing import IO, BinaryIO, TextIO, Tuple
+
+logger = logging.getLogger(__name__)
+
+
+LFS_VERSION = "https://git-lfs.github.com/spec/v1"
+LEGACY_LFS_VERSION = "https://hawser.github.com/spec/v1"
+ALLOWED_VERSIONS = (LFS_VERSION, LEGACY_LFS_VERSION)
+HEADERS = [(b"version " + version.encode("utf-8")) for version in ALLOWED_VERSIONS]
+
+
+def _get_kv(line: str) -> Tuple[str, str]:
+    key, value = line.strip().split(maxsplit=1)
+    return key, value
+
+
+@dataclass
+class Pointer:
+    oid: str
+    size: int
+
+    def __init__(self, oid: str, size: int, **kwargs):
+        self.oid = oid
+        self.size = size
+        self._dict = kwargs
+
+    def __hash__(self):
+        return hash(self.dump())
+
+    @classmethod
+    def build(cls, fobj: BinaryIO) -> "Pointer":
+        m = hashlib.sha256()
+        data = fobj.read()
+        m.update(data)
+        return cls(m.hexdigest(), len(data))
+
+    @classmethod
+    def load(cls, fobj: IO) -> "Pointer":
+        """Load the specified pointer file."""
+
+        if isinstance(fobj, io.TextIOBase):  # type: ignore[unreachable]
+            text_obj: TextIO = fobj  # type: ignore[unreachable]
+
+        else:
+            text_obj = io.TextIOWrapper(fobj, encoding="utf-8")
+
+        cls.check_version(text_obj.readline())
+        d = dict(_get_kv(line) for line in text_obj.readlines())
+        try:
+            value = d.pop("oid")
+            hash_method, oid = value.split(":", maxsplit=1)
+            if hash_method != "sha256":
+                raise ValueError("Invalid LFS hash method '{hash_method}'")
+        except ValueError as e:
+            raise ValueError("Invalid LFS pointer oid") from e
+        try:
+            value = d.pop("size")
+            size = int(value)
+        except ValueError as e:
+            raise ValueError("Invalid LFS pointer size") from e
+
+        return cls(oid, size, **d)
+
+    @staticmethod
+    def check_version(line: str):
+        try:
+            key, value = _get_kv(line)
+            if key != "version":
+                raise ValueError("LFS pointer file must start with 'version'")
+            if value not in ALLOWED_VERSIONS:
+                raise ValueError(f"Unsupported LFS pointer version '{value}'")
+        except (ValueError, OSError) as e:
+            raise ValueError("Invalid LFS pointer file") from e
+
+    def dump(self) -> str:
+        d = {
+            "oid": f"sha256:{self.oid}",
+            "size": self.size,
+        }
+        d.update(self._dict)
+        return "\n".join(
+            [f"version {LFS_VERSION}"] + [f"{key} {d[key]}" for key in sorted(d)] + [""]
+        )
+
+    def to_bytes(self) -> bytes:
+        return self.dump().encode("utf-8")
+
+
+if __name__ == "__main__":
+    # Minimal `git lfs pointer` CLI implementation
+    import argparse
+    import sys
+
+    parser = argparse.ArgumentParser(
+        description="Build generated pointer files.",
+    )
+    parser.add_argument("--file", help="A local file to build the pointer from.")
+    args = parser.parse_args()
+    if not args.file:
+        sys.exit("Nothing to do")
+
+    print(f"Git LFS pointer for {args.file}\n", file=sys.stderr)
+    with open(args.file, mode="rb") as fobj_:
+        p = Pointer.build(fobj_)
+    print(p.dump(), end="")

scmrepo/git/lfs/progress.py

@@ -0,0 +1,61 @@
+from typing import Any, BinaryIO, Callable, Dict, Optional, Union
+
+from dvc_objects.fs.callbacks import DEFAULT_CALLBACK, Callback, TqdmCallback
+
+from scmrepo.progress import GitProgressEvent
+
+
+class LFSCallback(Callback):
+    """Callback subclass to generate Git/LFS style progress."""
+
+    def __init__(
+        self,
+        *args,
+        git_progress: Optional[Callable[[GitProgressEvent], None]] = None,
+        direction: str = "Downloading",
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+        self.direction = direction
+        self.git_progress = git_progress
+
+    def call(self, *args, **kwargs):
+        super().call(*args, **kwargs)
+        self._update_git()
+
+    def _update_git(self):
+        if not self.git_progress:
+            return
+        event = GitProgressEvent(
+            phase=f"{self.direction} LFS objects",
+            completed=self.value,
+            total=self.size,
+        )
+        self.git_progress(event)
+
+    def branch(
+        self,
+        path_1: Union[str, BinaryIO],
+        path_2: str,
+        kwargs: Dict[str, Any],
+        child: Optional[Callback] = None,
+    ):
+        if child:
+            pass
+        elif self.git_progress:
+            child = TqdmCallback(
+                bytes=True, desc=path_1 if isinstance(path_1, str) else path_2
+            )
+        else:
+            child = DEFAULT_CALLBACK
+        return super().branch(path_1, path_2, kwargs, child=child)
+
+    @classmethod
+    def as_lfs_callback(
+        cls,
+        git_progress: Optional[Callable[[GitProgressEvent], None]] = None,
+        **kwargs,
+    ):
+        if git_progress is None:
+            return DEFAULT_CALLBACK
+        return cls(git_progress=git_progress, **kwargs)

scmrepo/git/lfs/smudge.py

@@ -0,0 +1,51 @@
+import io
+import logging
+from typing import TYPE_CHECKING, BinaryIO, Optional
+
+from .pointer import HEADERS, Pointer
+
+if TYPE_CHECKING:
+    from .storage import LFSStorage
+
+logger = logging.getLogger(__name__)
+
+
+def smudge(
+    storage: "LFSStorage", fobj: BinaryIO, url: Optional[str] = None
+) -> BinaryIO:
+    """Wrap the specified binary IO stream and run LFS smudge if necessary."""
+    reader = io.BufferedReader(fobj)  # type: ignore[arg-type]
+    data = reader.peek(100)
+    if any(data.startswith(header) for header in HEADERS):
+        # read the pointer data into memory since the raw stream is unseekable
+        # and we may need to return it in fallback case
+        data = reader.read()
+        lfs_obj: Optional[BinaryIO] = None
+        try:
+            pointer = Pointer.load(io.BytesIO(data))
+            lfs_obj = storage.open(pointer, mode="rb", fetch_url=url)
+        except (ValueError, OSError):
+            logger.warning("Could not open LFS object, falling back to raw pointer")
+        if lfs_obj:
+            return lfs_obj
+        return io.BytesIO(data)
+    return reader
+
+
+if __name__ == "__main__":
+    # Minimal `git lfs smudge` CLI implementation
+    import sys
+
+    from scmrepo.git import Git
+
+    if sys.stdin.isatty():
+        sys.exit(
+            "Cannot read from STDIN: "
+            "This command should be run by the Git 'smudge' filter"
+        )
+    scm = Git()
+    try:
+        with smudge(scm.lfs_storage, sys.stdin.buffer) as fobj_:
+            sys.stdout.buffer.write(fobj_.read())
+    finally:
+        scm.close()

scmrepo/git/lfs/storage.py

@@ -0,0 +1,74 @@
+import errno
+import os
+from typing import TYPE_CHECKING, BinaryIO, Callable, Collection, Optional, Union
+
+from .pointer import Pointer
+from .progress import LFSCallback
+
+if TYPE_CHECKING:
+    from scmrepo.git import Git
+    from scmrepo.progress import GitProgressEvent
+
+
+class LFSStorage:
+    def __init__(self, path: str):
+        self.path = path
+
+    def fetch(
+        self,
+        url: str,
+        objects: Collection[Pointer],
+        progress: Optional[Callable[["GitProgressEvent"], None]] = None,
+    ):
+        from .client import LFSClient
+
+        with LFSCallback.as_lfs_callback(progress) as cb:
+            cb.set_size(len(objects))
+            with LFSClient.from_git_url(url) as client:
+                client.download(self, objects, callback=cb)
+
+    def oid_to_path(self, oid: str):
+        return os.path.join(self.path, "objects", oid[0:2], oid[2:4], oid)
+
+    def exists(self, obj: Union[Pointer, str]):
+        oid = obj if isinstance(obj, str) else obj.oid
+        path = self.oid_to_path(oid)
+        return os.path.exists(path)
+
+    def open(
+        self,
+        obj: Union[Pointer, str],
+        fetch_url: Optional[str] = None,
+        **kwargs,
+    ) -> BinaryIO:
+        oid = obj if isinstance(obj, str) else obj.oid
+        path = self.oid_to_path(oid)
+        try:
+            return open(path, **kwargs)  # pylint: disable=unspecified-encoding
+        except FileNotFoundError:
+            if not fetch_url or not isinstance(obj, Pointer):
+                raise
+        try:
+            self.fetch(fetch_url, [obj])
+        except BaseException as exc:
+            raise FileNotFoundError(
+                errno.ENOENT, os.strerror(errno.ENOENT), path
+            ) from exc
+        return open(path, **kwargs)  # pylint: disable=unspecified-encoding
+
+    def close(self):
+        pass
+
+
+def get_storage_path(scm: "Git") -> str:
+    """Return the LFS storage directory for the specified repository."""
+
+    config = scm.get_config()
+    git_dir = scm._get_git_dir(scm.root_dir)  # pylint: disable=protected-access
+    try:
+        path = config.get(("lfs",), "storage")
+        if os.path.isabs(path):
+            return path
+    except KeyError:
+        path = "lfs"
+    return os.path.join(git_dir, path)