tgzr.snap 0.0.1__py3-none-any.whl

tgzr/snap/__init__.py

@@ -0,0 +1 @@
+

tgzr/snap/_version.py

@@ -0,0 +1,34 @@
+# file generated by setuptools-scm
+# don't change, don't track in version control
+
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    "version",
+    "version_tuple",
+    "__commit_id__",
+    "commit_id",
+]
+
+TYPE_CHECKING = False
+if TYPE_CHECKING:
+    from typing import Tuple
+    from typing import Union
+
+    VERSION_TUPLE = Tuple[Union[int, str], ...]
+    COMMIT_ID = Union[str, None]
+else:
+    VERSION_TUPLE = object
+    COMMIT_ID = object
+
+version: str
+__version__: str
+__version_tuple__: VERSION_TUPLE
+version_tuple: VERSION_TUPLE
+commit_id: COMMIT_ID
+__commit_id__: COMMIT_ID
+
+__version__ = version = '0.0.1'
+__version_tuple__ = version_tuple = (0, 0, 1)
+
+__commit_id__ = commit_id = None

tgzr/snap/base_repository.py

@@ -0,0 +1,115 @@
+from __future__ import annotations
+from typing import TypeVar, Generic, get_args, Any
+
+from abc import ABC, abstractmethod
+from pathlib import Path
+import dataclasses
+
+
+@dataclasses.dataclass
+class VersionInfo:
+    name: str
+    isdir: bool = False
+    size: int | None = None
+    nfiles: int | None = None
+    isexec: bool = False
+    version_id: str | None = None
+    md5: str | None = None
+    mtime: float | None = None
+
+    remote_name: str | None = None
+
+
+@dataclasses.dataclass
+class Remote:
+    name: str
+    url: str
+
+
+@dataclasses.dataclass
+class RepositoryConfig:
+    remotes: list[Remote]
+    default_remote: str
+
+    @classmethod
+    def from_dict(cls, d: dict[str, Any]):
+        return cls(**d)
+
+
+RepositoryConfigType = TypeVar("RepositoryConfigType", bound=RepositoryConfig)
+
+
+class SnapRepository(ABC, Generic[RepositoryConfigType]):
+    @classmethod
+    @abstractmethod
+    def create_repository(
+        cls, path: str | Path, config: RepositoryConfig, force: bool = False
+    ):
+        """
+        Initialize a new repository a that path with this config.
+        """
+        pass
+
+    @classmethod
+    def get_config_type(cls) -> RepositoryConfigType:
+        return get_args(cls.__orig_bases__[0])[0]  # type: ignore __orig_bases__ trust me bro.
+
+    def __init__(self, path: str | Path, config: RepositoryConfigType):
+        self._path = Path(path)
+        self._config = config
+
+    @property
+    def config(self) -> RepositoryConfigType:
+        return self._config
+
+    @property
+    def path(self) -> Path:
+        return self._path
+
+    #
+    #
+    #
+
+    @abstractmethod
+    def default_remote_name(self) -> str:
+        pass
+
+    @abstractmethod
+    def remotes(self) -> list[Remote]:
+        pass
+
+    @abstractmethod
+    def get_remote(self, remote_name: str | None) -> Remote:
+        """
+        Raise KeyError if no such remote is declared.
+        """
+        pass
+
+    #
+    #
+    #
+
+    @abstractmethod
+    def push(self, repo_path: str, remote_name: str | None = None) -> VersionInfo:
+        """Uploads the file at `self.path/repo_path` and returns the VersionInfo for the created version."""
+        pass
+
+    @abstractmethod
+    def get_versions(self, repo_path: str):
+        """Returns a list of VersionInfo for repo_path."""
+        pass
+
+    @abstractmethod
+    def status(self, repo_path: str):
+        """Returns the status of the file/folder at repo_path."""
+        pass
+
+    @abstractmethod
+    def pull(self, repo_path: str | Path, version: str, remote_name: str | None = None):
+        """Downloads this version of `repo_path`."""
+        pass
+
+    @abstractmethod
+    def unprotect(self, repo_path: str):
+        """Make the file writable."""
+        pass

tgzr/snap/plugin.py

@@ -0,0 +1,36 @@
+from __future__ import annotations
+from typing import Type
+
+from tgzr.package_management.plugin_manager import Plugin, PluginManager
+
+from .base_repository import SnapRepository
+
+
+class SnapPlugin(Plugin):
+    @classmethod
+    def plugin_type_name(cls) -> str:
+        return "SnapRepository"
+
+    def get_repository_type(self) -> Type[SnapRepository]:
+        raise NotImplementedError()
+
+
+class SnapPluginManager(PluginManager[SnapPlugin]):
+    EP_GROUP = "tgzr.snap.plugin"
+
+
+def test_plugins():
+    # NB: we can't use the manager from here or the SnapPlugin qualified name
+    # is '__main__.SnapPluginb' instead of 'tgzr.snap.plugin.SnapPlugin'.
+    # so:
+    from tgzr.snap.plugin import SnapPluginManager
+
+    pm = SnapPluginManager()
+    print("Snap Plugins:")
+    for plugin in pm.get_plugins():
+        snap_type = plugin.get_repository_type()
+        print("   ", snap_type, snap_type.get_config_type())
+
+
+if __name__ == "__main__":
+    test_plugins()

tgzr/snap/plugins/__init__.py

@@ -0,0 +1,19 @@
+from __future__ import annotations
+from typing import Type
+
+from .dvc import DVCSnap
+from .folder import FolderSnap
+
+from ..plugin import SnapPlugin, SnapRepository
+
+
+class DVCSnapPlugin(SnapPlugin):
+
+    def get_repository_type(self) -> Type[SnapRepository]:
+        return DVCSnap
+
+
+class FolderSnapPlugin(SnapPlugin):
+
+    def get_repository_type(self) -> Type[SnapRepository]:
+        return FolderSnap

tgzr/snap/plugins/dvc.py

@@ -0,0 +1,363 @@
+from __future__ import annotations
+
+import os
+import shutil
+from pathlib import Path
+import dataclasses
+
+import rich
+import json
+
+from dvc.repo import Repo
+from dvc_data.hashfile.hash_info import HashInfo
+from dvc_data.hashfile.meta import Meta
+from dvc_data.hashfile.build import build
+from dvc_data.hashfile.tree import Tree
+
+from dvc_data.hashfile.build import _build_tree
+from dvc_data.hashfile.db import add_update_tree
+
+from dvc_data.hashfile.hash import hash_file
+from dvc_data.hashfile.db import HashFileDB
+from dvc_data.hashfile.transfer import transfer
+from dvc_data.hashfile.obj import HashFile
+from dvc_data.hashfile.checkout import checkout
+
+
+from dvc.stage import Stage
+from dvc.output import Output
+
+from ..base_repository import SnapRepository, RepositoryConfig, VersionInfo, Remote
+
+
+@dataclasses.dataclass
+class DVCRepositoryConfig(RepositoryConfig):
+    cache_dir: str | Path
+
+
+class DVCSnap(SnapRepository[DVCRepositoryConfig]):
+
+    @classmethod
+    def create_repository(
+        cls, path: str | Path, config: DVCRepositoryConfig, force: bool = False
+    ) -> DVCSnap:
+        """
+        Create the dvc repository on disk and return a DVCSnap connected to it.
+        The remotes defined in config will be configured in the dvc repository.
+        """
+        repo = Repo.init(str(path), no_scm=True, force=force)
+        level = "local"
+        with repo.config.edit(level=level) as repo_config:
+            repo_config["core"]["analytics"] = False
+            repo_config["core"]["hardlink_lock"] = False
+            repo_config["core"]["check_update"] = False
+
+            repo_config["cache"]["dir"] = str(config.cache_dir)
+            repo_config["cache"]["type"] = "reflink,hardlink,copy"
+            repo_config["cache"]["protected"] = True
+
+            if "remote" not in repo_config:
+                repo_config["remote"] = {}
+
+            remote_names = []
+            for remote in config.remotes:
+                repo_config["remote"][remote.name] = {"url": remote.url}
+                remote_names.append(remote.name)
+
+            if config.default_remote:
+                if not config.default_remote in remote_names:
+                    raise ValueError(
+                        f"Cannot use remote {config.default_remote} as defaut: no such remote in config (got {remote_names})"
+                    )
+                repo_config["core"]["remote"] = config.default_remote
+
+        # rich.print(repo.config)
+
+        dvc_snap = cls(path, config)
+        return dvc_snap
+
+    def __init__(self, path: str | Path, config: DVCRepositoryConfig):
+        super().__init__(path, config)
+
+        self._dvc_repo = Repo(str(self.path))
+
+    #
+    #
+    #
+
+    def default_remote_name(self) -> str:
+        rich.print(self._dvc_repo.config)
+        return self._dvc_repo.config.get("core", {}).get("remote")
+
+    def remotes(self) -> list[Remote]:
+        remotes_config = self._dvc_repo.config.get("remote", {})
+
+        return [
+            Remote(name, options.get("url", None))
+            for name, options in remotes_config.items()
+        ]
+
+    def get_remote(self, remote_name: str | None) -> Remote:
+        remote_name = remote_name or self.default_remote_name()
+        r = self._dvc_repo.config["remote"][remote_name]
+        return Remote(name=remote_name, url=r["url"])
+
+    #
+    #
+    #
+
+    def push(
+        self, repo_path: str | Path, remote_name: str | None = None
+    ) -> VersionInfo:
+        path = str(self.path / repo_path)
+
+        remote = self.get_remote(remote_name)
+        local_odb = self._dvc_repo.cache.local
+        remote_odb = self._dvc_repo.cloud.get_remote_odb(remote.name)
+
+        _, meta, tree_or_hashfile_obj = build(local_odb, path, local_odb.fs, "md5")
+
+        if tree_or_hashfile_obj.hash_info.isdir:
+            for entry_path, entry_meta, entry_hi in tree_or_hashfile_obj:
+                if not local_odb.exists(entry_hi.value):
+                    # Resolve the full path for the child file
+                    full_entry_path = local_odb.fs.join(path, *entry_path)
+                    local_odb.add(full_entry_path, local_odb.fs, entry_hi.value)
+            add_update_tree(local_odb, tree_or_hashfile_obj)
+
+            # it is a Tree object (Pushing a folder)
+            to_push = {tree_or_hashfile_obj.hash_info}
+            for _, _, child_hi in tree_or_hashfile_obj:
+                to_push.add(child_hi)
+        else:
+            local_odb.add(path, local_odb.fs, tree_or_hashfile_obj.hash_info.value)
+            to_push = {tree_or_hashfile_obj.hash_info}
+
+        transfer(
+            local_odb,
+            remote_odb,
+            to_push,
+            jobs=None,
+        )
+        print("--->", tree_or_hashfile_obj.hash_info, meta)
+        assert (
+            tree_or_hashfile_obj.hash_info.value is not None
+        )  # should alway be true, makes mypy happy.
+
+        vi = VersionInfo(
+            name=tree_or_hashfile_obj.hash_info.value,
+            isdir=meta.isdir,
+            size=meta.size,
+            nfiles=meta.nfiles,
+            isexec=meta.isexec,
+            version_id=meta.version_id,
+            md5=meta.md5,
+            mtime=meta.mtime,
+            remote_name=remote_name or self.default_remote_name(),
+        )
+        return vi
+
+    def get_versions(self, repo_path: str):
+        path = self.path / repo_path
+        raise NotImplementedError()
+
+    def status(self, repo_path: str):
+        path = self.path / repo_path
+        raise NotImplementedError()
+
+    def pull(self, repo_path: str | Path, version: str, remote_name: str | None = None):
+        path = str(self.path / repo_path)
+
+        local_odb = self._dvc_repo.cache.local
+        remote_odb = self._dvc_repo.cloud.get_remote_odb(remote_name)
+
+        hash_info = HashInfo("md5", version)
+
+        if not local_odb.exists(hash_info.value):
+            transfer(remote_odb, local_odb, {hash_info})
+
+        obj = local_odb.get(hash_info.value)
+
+        if hash_info.isdir:
+            # tree_obj = Tree.from_dict(obj.hash_info, obj.fs, obj.path)
+            tree_obj = Tree.load(local_odb, obj.hash_info)
+            to_pull = {hash_info}
+            for _, _, child_hi in tree_obj:
+                to_pull.add(child_hi)
+            transfer(remote_odb, local_odb, to_pull)
+            checkout_obj = tree_obj
+        else:
+            checkout_obj = local_odb.get(hash_info.value)
+
+        checkout(
+            path,
+            local_odb.fs,
+            checkout_obj,
+            local_odb,
+            force=True,
+            state=self._dvc_repo.state,
+        )
+
+        if self._dvc_repo.state:
+            self._dvc_repo.state.save(path, local_odb.fs, hash_info)
+
+    def unprotect(self, repo_path: str):
+        path = self.path / repo_path
+        stage = Stage(self._dvc_repo)
+        out = Output(stage, str(path))
+        out.unprotect()
+
+
+#
+#
+#
+
+
+def test_config() -> DVCRepositoryConfig:
+    config = DVCRepositoryConfig(
+        remotes=[
+            Remote("blessed", "/tmp/snap_tests_dvc/PROD"),
+            Remote("review", "/tmp/snap_tests_dvc/REVIEW"),
+        ],
+        default_remote="review",
+        cache_dir="/tmp/snap_tests_dvc/CACHE",
+    )
+    return config
+
+
+def create_repo():
+    """
+    mkdir PROD_REPO
+
+    mkdir my_data_repo && cd my_data_repo
+    dvc init --no-scm
+    dvc config cache.type "reflink,hardlink,copy"
+    dvc config cache.protected true
+    dvc remote add -d PROD ../PROD_REPO
+
+    echo "content" > bob.txt
+    """
+    config = test_config()
+    path = Path("/tmp/snap_tests_dvc/WORK")
+    force = True
+    DVCSnap.create_repository(path, config, force=force)
+
+    # snap = DVCSnap(path, config)
+    # rich.print(snap._dvc_repo.config)
+
+
+def test_push():
+    import time
+
+    config = test_config()
+    path = Path("/tmp/snap_tests_dvc/WORK")
+    snap = DVCSnap(path, config)
+    artifact = "bob.txt"
+
+    snap.unprotect(artifact)
+
+    with open(snap.path / artifact, "w") as fp:
+        fp.write(f"Change... {time.time()}")
+
+    version = snap.push(artifact, "blessed")
+
+    print(f"{version=}")
+
+
+def test_pull():
+
+    config = test_config()
+    path = Path("/tmp/snap_tests_dvc/WORK")
+    snap = DVCSnap(path, config)
+    artifact = "bob.txt"
+
+    version = "ccef8ce5e5b64d56cc72ef9e118df82d"
+    version = "b1ef6fd0c3d7fb7ebf4769593b9c445b"
+    version = "ac6281fd84ac444e59af60f53da81a48"
+
+    snap.pull(artifact, version)
+    with open(path / artifact, "r") as fp:
+        print("CONTENT is:", fp.read())
+
+
+def test_push_iso():
+    config = test_config()
+    path = Path("/tmp/snap_tests_dvc/WORK")
+    snap = DVCSnap(path, config)
+    artifact = "big.iso"
+
+    Mo = 360000
+    with open(path / artifact, "w") as fp:
+        fp.write(str(os.urandom(100 * Mo)))
+
+    version_info = snap.push(artifact)
+
+    rich.print(f"{version_info}")
+    print(f"    version='{version_info.name}'")
+
+
+def test_pull_iso():
+    import time
+
+    config = test_config()
+    path = Path("/tmp/snap_tests_dvc/WORK")
+    snap = DVCSnap(path, config)
+    artifact = "big.iso"
+
+    version = "1b9cb5c1940958f28552c5ce7be1a3bc"  # 928K
+    version = "8f33876d8bc03223ff4bcb70b772dec7"  # 1010K
+    version = "174101e0ee8b63f760d5994b4c2e6380"  # 99M
+
+    t = time.time()
+    snap.pull(artifact, version)
+    print(time.time() - t)
+
+
+def test_push_dir():
+    import random
+
+    config = test_config()
+    path = Path("/tmp/snap_tests_dvc/WORK")
+    snap = DVCSnap(path, config)
+    artifact = "data_folder"
+
+    (path / artifact).mkdir(exist_ok=True)
+
+    for i in range(random.randrange(5, 10)):
+        Mo = 360000
+        size = random.randrange(1, 20)
+        with open(path / artifact / f"file_{i:03}.data", "w") as fp:
+            fp.write(str(os.urandom(size * Mo)))
+
+    version_info = snap.push(artifact)
+
+    rich.print(f"{version_info}")
+    print(f"    version='{version_info.name}'")
+
+
+def test_pull_dir():
+
+    import time
+
+    config = test_config()
+    path = Path("/tmp/snap_tests_dvc/WORK")
+    snap = DVCSnap(path, config)
+    artifact = "data_folder"
+
+    version = "225b0fbdaf1dbeaa3f0dd7ffeecd716f.dir"
+    version = "4ca4e81d1533353d2fe7ed1bc801d39f.dir"
+    version = "fd0c0e8ab5c798c61778f63063acc462.dir"
+
+    t = time.time()
+    snap.pull(artifact, version)
+    print(time.time() - t)
+
+
+if __name__ == "__main__":
+    # create_repo()
+    # test_push()
+    # test_pull()
+    # test_push_iso()
+    # test_pull_iso()
+    # test_push_dir()
+    test_pull_dir()