objects: split into separate files (#5609)

dvc/checkout.py

@@ -9,7 +9,8 @@ from dvc.exceptions import (
     ConfirmRemoveError,
     DvcException,
 )
-from dvc.objects import ObjectFormatError, check, load
+from dvc.objects import check, load
+from dvc.objects.errors import ObjectFormatError
 from dvc.objects.stage import stage
 from dvc.remote.slow_link_detection import (  # type: ignore[attr-defined]
     slow_link_guard,

dvc/objects/__init__.py

@@ -1,77 +1,13 @@
-import itertools
 import logging
-from concurrent import futures
 
-from dvc.exceptions import DvcException
 from dvc.progress import Tqdm
 
-from .stage import get_file_hash
+from .tree import Tree
 
 logger = logging.getLogger(__name__)
 
 
-class ObjectError(DvcException):
-    pass
-
-
-class ObjectFormatError(ObjectError):
-    pass
-
-
-class HashFile:
-    def __init__(self, path_info, fs, hash_info):
-        self.path_info = path_info
-        self.fs = fs
-        self.hash_info = hash_info
-
-    @property
-    def size(self):
-        if not (self.path_info and self.fs):
-            return None
-        return self.fs.getsize(self.path_info)
-
-    def __len__(self):
-        return 1
-
-    def __str__(self):
-        return f"object {self.hash_info}"
-
-    def __bool__(self):
-        return bool(self.hash_info)
-
-    def __eq__(self, other):
-        if not isinstance(other, HashFile):
-            return False
-        return (
-            self.path_info == other.path_info
-            and self.fs == other.fs
-            and self.hash_info == other.hash_info
-        )
-
-    def check(self, odb):
-        actual = get_file_hash(
-            self.path_info, self.fs, self.hash_info.name, odb.repo.state
-        )
-
-        logger.trace(
-            "cache '%s' expected '%s' actual '%s'",
-            self.path_info,
-            self.hash_info,
-            actual,
-        )
-
-        assert actual.name == self.hash_info.name
-        if actual.value.split(".")[0] != self.hash_info.value.split(".")[0]:
-            raise ObjectFormatError(f"{self} is corrupted")
-
-    @classmethod
-    def load(cls, odb, hash_info):
-        return odb.get(hash_info)
-
-
 def save(odb, obj, **kwargs):
-    from .tree import Tree
-
     if isinstance(obj, Tree):
         for _, entry in Tqdm(obj):
             odb.add(entry.path_info, entry.fs, entry.hash_info, **kwargs)
@@ -79,8 +15,6 @@ def save(odb, obj, **kwargs):
 
 
 def check(odb, obj):
-    from .tree import Tree
-
     odb.check(obj.hash_info)
 
     if isinstance(obj, Tree):
@@ -89,91 +23,6 @@ def check(odb, obj):
 
 
 def load(odb, hash_info):
-    from .tree import Tree
-
     if hash_info.isdir:
         return Tree.load(odb, hash_info)
     return odb.get(hash_info)
-
-
-def _transfer_file(odb, from_fs, from_info):
-    from dvc.utils import tmp_fname
-    from dvc.utils.stream import HashedStreamReader
-
-    tmp_info = odb.fs.path_info / tmp_fname()
-    with from_fs.open(
-        from_info, mode="rb", chunk_size=from_fs.CHUNK_SIZE
-    ) as stream:
-        stream_reader = HashedStreamReader(stream)
-        # Since we don't know the hash beforehand, we'll
-        # upload it to a temporary location and then move
-        # it.
-        odb.fs.upload_fobj(
-            stream_reader,
-            tmp_info,
-            total=from_fs.getsize(from_info),
-            desc=from_info.name,
-        )
-
-    hash_info = stream_reader.hash_info
-    return tmp_info, hash_info
-
-
-def _transfer_directory_contents(odb, from_fs, from_info, jobs, pbar):
-    rel_path_infos = {}
-    from_infos = from_fs.walk_files(from_info)
-
-    def create_tasks(executor, amount):
-        for entry_info in itertools.islice(from_infos, amount):
-            pbar.total += 1
-            task = executor.submit(
-                pbar.wrap_fn(_transfer_file), odb, from_fs, entry_info
-            )
-            rel_path_infos[task] = entry_info.relative_to(from_info)
-            yield task
-
-    pbar.total = 0
-    with futures.ThreadPoolExecutor(max_workers=jobs) as executor:
-        tasks = set(create_tasks(executor, jobs * 5))
-
-        while tasks:
-            done, tasks = futures.wait(
-                tasks, return_when=futures.FIRST_COMPLETED
-            )
-            tasks.update(create_tasks(executor, len(done)))
-            for task in done:
-                yield rel_path_infos.pop(task), task.result()
-
-
-def _transfer_directory(odb, from_fs, from_info, jobs, no_progress_bar=False):
-    from .tree import Tree
-
-    tree = Tree(None, None, None)
-
-    with Tqdm(total=1, unit="Files", disable=no_progress_bar) as pbar:
-        for (
-            entry_info,
-            (entry_tmp_info, entry_hash),
-        ) in _transfer_directory_contents(odb, from_fs, from_info, jobs, pbar):
-            obj = HashFile(entry_tmp_info, odb.fs, entry_hash)
-            tree.add(entry_info.parts, obj)
-
-    tree.digest()
-    save(odb, tree)
-    return tree.hash_info
-
-
-def transfer(odb, from_fs, from_info, jobs=None, no_progress_bar=False):
-    jobs = jobs or min((from_fs.jobs, odb.fs.jobs))
-
-    if from_fs.isdir(from_info):
-        return _transfer_directory(
-            odb,
-            from_fs,
-            from_info,
-            jobs=jobs,
-            no_progress_bar=no_progress_bar,
-        )
-    tmp_info, hash_info = _transfer_file(odb, from_fs, from_info)
-    odb.add(tmp_info, odb.fs, hash_info)
-    return hash_info

dvc/objects/db/base.py

@@ -4,10 +4,10 @@ from concurrent.futures import ThreadPoolExecutor
 from copy import copy
 from typing import Optional
 
+from dvc.objects.errors import ObjectFormatError
+from dvc.objects.file import HashFile
 from dvc.progress import Tqdm
 
-from .. import HashFile, ObjectFormatError
-
 logger = logging.getLogger(__name__)
 
 

dvc/objects/db/local.py

@@ -6,7 +6,7 @@ from funcy import cached_property
 from shortuuid import uuid
 
 from dvc.hash_info import HashInfo
-from dvc.objects import ObjectFormatError
+from dvc.objects.errors import ObjectFormatError
 from dvc.path_info import PathInfo
 from dvc.progress import Tqdm
 from dvc.utils import relpath

dvc/objects/errors.py

@@ -0,0 +1,9 @@
+from dvc.exceptions import DvcException
+
+
+class ObjectError(DvcException):
+    pass
+
+
+class ObjectFormatError(ObjectError):
+    pass

dvc/objects/file.py

@@ -0,0 +1,57 @@
+import logging
+
+from .errors import ObjectFormatError
+from .stage import get_file_hash
+
+logger = logging.getLogger(__name__)
+
+
+class HashFile:
+    def __init__(self, path_info, fs, hash_info):
+        self.path_info = path_info
+        self.fs = fs
+        self.hash_info = hash_info
+
+    @property
+    def size(self):
+        if not (self.path_info and self.fs):
+            return None
+        return self.fs.getsize(self.path_info)
+
+    def __len__(self):
+        return 1
+
+    def __str__(self):
+        return f"object {self.hash_info}"
+
+    def __bool__(self):
+        return bool(self.hash_info)
+
+    def __eq__(self, other):
+        if not isinstance(other, HashFile):
+            return False
+        return (
+            self.path_info == other.path_info
+            and self.fs == other.fs
+            and self.hash_info == other.hash_info
+        )
+
+    def check(self, odb):
+        actual = get_file_hash(
+            self.path_info, self.fs, self.hash_info.name, odb.repo.state
+        )
+
+        logger.trace(
+            "cache '%s' expected '%s' actual '%s'",
+            self.path_info,
+            self.hash_info,
+            actual,
+        )
+
+        assert actual.name == self.hash_info.name
+        if actual.value.split(".")[0] != self.hash_info.value.split(".")[0]:
+            raise ObjectFormatError(f"{self} is corrupted")
+
+    @classmethod
+    def load(cls, odb, hash_info):
+        return odb.get(hash_info)

dvc/objects/stage.py

@@ -77,7 +77,7 @@ def _iter_hashes(path_info, fs, name, state, **kwargs):
 
 
 def _build_tree(path_info, fs, name, state, **kwargs):
-    from . import HashFile
+    from .file import HashFile
     from .tree import Tree
 
     tree = Tree(None, None, None)
@@ -176,7 +176,7 @@ def stage(odb, path_info, fs, name, **kwargs):
     if fs.isdir(path_info):
         obj = _get_tree_obj(path_info, fs, name, odb, state, **kwargs)
     else:
-        from . import HashFile
+        from .file import HashFile
 
         hash_info = get_file_hash(path_info, fs, name, state)
         obj = HashFile(path_info, fs, hash_info)

dvc/objects/transfer.py

@@ -0,0 +1,92 @@
+import itertools
+import logging
+from concurrent import futures
+
+from dvc.progress import Tqdm
+
+from . import save
+from .file import HashFile
+from .tree import Tree
+
+logger = logging.getLogger(__name__)
+
+
+def _transfer_file(odb, from_fs, from_info):
+    from dvc.utils import tmp_fname
+    from dvc.utils.stream import HashedStreamReader
+
+    tmp_info = odb.fs.path_info / tmp_fname()
+    with from_fs.open(
+        from_info, mode="rb", chunk_size=from_fs.CHUNK_SIZE
+    ) as stream:
+        stream_reader = HashedStreamReader(stream)
+        # Since we don't know the hash beforehand, we'll
+        # upload it to a temporary location and then move
+        # it.
+        odb.fs.upload_fobj(
+            stream_reader,
+            tmp_info,
+            total=from_fs.getsize(from_info),
+            desc=from_info.name,
+        )
+
+    hash_info = stream_reader.hash_info
+    return tmp_info, hash_info
+
+
+def _transfer_directory_contents(odb, from_fs, from_info, jobs, pbar):
+    rel_path_infos = {}
+    from_infos = from_fs.walk_files(from_info)
+
+    def create_tasks(executor, amount):
+        for entry_info in itertools.islice(from_infos, amount):
+            pbar.total += 1
+            task = executor.submit(
+                pbar.wrap_fn(_transfer_file), odb, from_fs, entry_info
+            )
+            rel_path_infos[task] = entry_info.relative_to(from_info)
+            yield task
+
+    pbar.total = 0
+    with futures.ThreadPoolExecutor(max_workers=jobs) as executor:
+        tasks = set(create_tasks(executor, jobs * 5))
+
+        while tasks:
+            done, tasks = futures.wait(
+                tasks, return_when=futures.FIRST_COMPLETED
+            )
+            tasks.update(create_tasks(executor, len(done)))
+            for task in done:
+                yield rel_path_infos.pop(task), task.result()
+
+
+def _transfer_directory(odb, from_fs, from_info, jobs, no_progress_bar=False):
+    tree = Tree(None, None, None)
+
+    with Tqdm(total=1, unit="Files", disable=no_progress_bar) as pbar:
+        for (
+            entry_info,
+            (entry_tmp_info, entry_hash),
+        ) in _transfer_directory_contents(odb, from_fs, from_info, jobs, pbar):
+            obj = HashFile(entry_tmp_info, odb.fs, entry_hash)
+            tree.add(entry_info.parts, obj)
+
+    tree.digest()
+    save(odb, tree)
+    return tree.hash_info
+
+
+def transfer(odb, from_fs, from_info, jobs=None, no_progress_bar=False):
+    jobs = jobs or min((from_fs.jobs, odb.fs.jobs))
+
+    if from_fs.isdir(from_info):
+        return _transfer_directory(
+            odb,
+            from_fs,
+            from_info,
+            jobs=jobs,
+            no_progress_bar=no_progress_bar,
+        )
+    tmp_info, hash_info = _transfer_file(odb, from_fs, from_info)
+    odb.add(tmp_info, odb.fs, hash_info)
+    return hash_info

dvc/objects/tree.py

@@ -4,7 +4,8 @@ import posixpath
 
 from funcy import cached_property
 
-from . import HashFile, ObjectFormatError
+from .errors import ObjectFormatError
+from .file import HashFile
 from .stage import get_file_hash
 
 logger = logging.getLogger(__name__)

dvc/output/base.py

@@ -18,6 +18,7 @@ from dvc.exceptions import (
 )
 from dvc.hash_info import HashInfo
 from dvc.objects.db import NamedCache
+from dvc.objects.errors import ObjectFormatError
 from dvc.objects.stage import stage as ostage
 
 from ..fs.base import BaseFileSystem
@@ -220,7 +221,7 @@ class BaseOutput:
         try:
             objects.check(self.odb, obj)
             return False
-        except (FileNotFoundError, objects.ObjectFormatError):
+        except (FileNotFoundError, ObjectFormatError):
             return True
 
     def workspace_status(self):
@@ -487,7 +488,7 @@ class BaseOutput:
 
         try:
             objects.check(self.odb, self.odb.get(self.hash_info))
-        except (FileNotFoundError, objects.ObjectFormatError):
+        except (FileNotFoundError, ObjectFormatError):
             self.repo.cloud.pull(
                 NamedCache.make("local", self.hash_info.value, str(self)),
                 show_checksums=False,
@@ -496,7 +497,7 @@ class BaseOutput:
 
         try:
             self.obj = objects.load(self.odb, self.hash_info)
-        except (objects.ObjectFormatError, FileNotFoundError):
+        except (FileNotFoundError, ObjectFormatError):
             self.obj = None
 
         return self.obj
@@ -532,7 +533,7 @@ class BaseOutput:
 
         try:
             objects.check(self.odb, self.odb.get(self.hash_info))
-        except (FileNotFoundError, objects.ObjectFormatError):
+        except (FileNotFoundError, ObjectFormatError):
             msg = (
                 "Missing cache for directory '{}'. "
                 "Cache for files inside will be lost. "

dvc/remote/base.py

@@ -506,7 +506,7 @@ class Remote:
         return ret
 
     def transfer(self, from_fs, from_info, jobs=None, no_progress_bar=False):
-        from dvc.objects import transfer
+        from dvc.objects.transfer import transfer
 
         return transfer(
             self.odb,

dvc/repo/add.py

@@ -173,7 +173,7 @@ def _process_stages(
             )
         else:
             from dvc.fs import get_cloud_fs
-            from dvc.objects import transfer
+            from dvc.objects.transfer import transfer
 
             from_fs = get_cloud_fs(repo, url=target)
             out.hash_info = transfer(

tests/func/test_odb.py

@@ -7,8 +7,8 @@ import pytest
 
 from dvc.hash_info import HashInfo
 from dvc.main import main
-from dvc.objects import ObjectFormatError
 from dvc.objects.db import ODBManager
+from dvc.objects.errors import ObjectFormatError
 from dvc.utils import relpath
 from tests.basic_env import TestDir, TestDvc
 

tests/unit/objects/test_tree.py

@@ -3,7 +3,7 @@ from operator import itemgetter
 import pytest
 
 from dvc.hash_info import HashInfo
-from dvc.objects import HashFile
+from dvc.objects.file import HashFile
 from dvc.objects.tree import Tree, _merge