objects: separate staging (#5470)

dvc/api/__init__.py

@@ -4,7 +4,7 @@ from contextlib import _GeneratorContextManager as GCM
 from funcy import reraise
 
 from dvc.exceptions import OutputNotFoundError, PathMissingError
-from dvc.oid import get_hash
+from dvc.objects.stage import get_hash
 from dvc.path_info import PathInfo
 from dvc.repo import Repo
 

dvc/checkout.py

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

dvc/dependency/repo.py

@@ -47,7 +47,7 @@ class RepoDependency(LocalDependency):
         return external_repo(d["url"], rev=rev, **kwargs)
 
     def _get_hash(self, locked=True):
-        from dvc.oid import get_hash
+        from dvc.objects.stage import get_hash
 
         with self._make_repo(locked=locked) as repo:
             path_info = PathInfo(repo.root_dir) / self.def_path
@@ -75,7 +75,8 @@ class RepoDependency(LocalDependency):
 
     def download(self, to, jobs=None):
         from dvc.checkout import checkout
-        from dvc.objects import save, stage
+        from dvc.objects import save
+        from dvc.objects.stage import stage
 
         cache = self.repo.cache.local
 

dvc/objects/__init__.py

@@ -3,11 +3,11 @@ import json
 import logging
 from concurrent import futures
 
+from dvc.dir_info import DirInfo
+from dvc.exceptions import DvcException
 from dvc.progress import Tqdm
 
-from .dir_info import DirInfo
-from .exceptions import DvcException
-from .oid import get_hash
+from .stage import get_hash
 
 logger = logging.getLogger(__name__)
 
@@ -189,12 +189,6 @@ class Tree(HashFile):
         return load(odb, self.save_dir_info(odb, dir_info))
 
 
-def stage(odb, path_info, fs, **kwargs):
-    if fs.isdir(path_info):
-        return Tree.stage(odb, path_info, fs, **kwargs)
-    return File.stage(odb, path_info, fs, **kwargs)
-
-
 def save(odb, obj, **kwargs):
     obj.save(odb, **kwargs)
 

dvc/objects/stage.py

@@ -89,7 +89,7 @@ def _collect_dir(path_info, fs, name, **kwargs):
 
 
 def get_dir_hash(path_info, fs, name, **kwargs):
-    from dvc.objects import Tree
+    from . import Tree
 
     dir_info = _collect_dir(path_info, fs, name, **kwargs)
     hash_info = Tree.save_dir_info(fs.repo.cache.local, dir_info)
@@ -126,7 +126,7 @@ def get_hash(path_info, fs, name, **kwargs):
 
         if hash_info:
             if hash_info.isdir:
-                from dvc.objects import Tree
+                from . import Tree
 
                 # NOTE: loading the fs will restore hash_info.dir_info
                 Tree.load(fs.cache, hash_info)
@@ -142,3 +142,11 @@ def get_hash(path_info, fs, name, **kwargs):
             fs.state.save(path_info, hash_info)
 
     return hash_info
+
+
+def stage(odb, path_info, fs, **kwargs):
+    from . import File, Tree
+
+    if fs.isdir(path_info):
+        return Tree.stage(odb, path_info, fs, **kwargs)
+    return File.stage(odb, path_info, fs, **kwargs)

dvc/output/base.py

@@ -18,7 +18,8 @@ from dvc.exceptions import (
     RemoteCacheRequiredError,
 )
 from dvc.hash_info import HashInfo
-from dvc.oid import get_hash
+from dvc.objects.stage import get_hash
+from dvc.objects.stage import stage as ostage
 
 from ..fs.base import BaseFileSystem
 
@@ -193,7 +194,7 @@ class BaseOutput:
     def get_hash(self):
         if not self.use_cache:
             return get_hash(self.path_info, self.fs, self.fs.PARAM_CHECKSUM)
-        return objects.stage(self.cache, self.path_info, self.fs).hash_info
+        return ostage(self.cache, self.path_info, self.fs).hash_info
 
     @property
     def is_dir_checksum(self):
@@ -301,7 +302,7 @@ class BaseOutput:
             logger.debug("Output '%s' didn't change. Skipping saving.", self)
             return
 
-        self.obj = objects.stage(self.cache, self.path_info, self.fs)
+        self.obj = ostage(self.cache, self.path_info, self.fs)
         self.hash_info = self.obj.hash_info
         self.isexec = self.isfile() and self.fs.isexec(self.path_info)
 
@@ -316,9 +317,7 @@ class BaseOutput:
         assert self.hash_info
 
         if self.use_cache:
-            obj = objects.stage(
-                self.cache, filter_info or self.path_info, self.fs
-            )
+            obj = ostage(self.cache, filter_info or self.path_info, self.fs)
             objects.save(self.cache, obj)
             checkout(
                 filter_info or self.path_info,

dvc/repo/diff.py

@@ -2,7 +2,7 @@ import logging
 import os
 
 from dvc.exceptions import PathMissingError
-from dvc.oid import get_file_hash, get_hash
+from dvc.objects.stage import get_file_hash, get_hash
 from dvc.repo import locked
 
 logger = logging.getLogger(__name__)

dvc/repo/fetch.py

@@ -81,7 +81,8 @@ def fetch(
 
 def _fetch_external(self, repo_url, repo_rev, files, jobs):
     from dvc.external_repo import external_repo
-    from dvc.objects import save, stage
+    from dvc.objects import save
+    from dvc.objects.stage import stage
     from dvc.path_info import PathInfo
     from dvc.scm.base import CloneError
 

tests/func/test_add.py

@@ -443,7 +443,7 @@ class TestDoubleAddUnchanged(TestDvc):
 
 
 def test_should_update_state_entry_for_file_after_add(mocker, dvc, tmp_dir):
-    file_md5_counter = mocker.spy(dvc_module.oid, "file_md5")
+    file_md5_counter = mocker.spy(dvc_module.objects.stage, "file_md5")
     tmp_dir.gen("foo", "foo")
 
     ret = main(["config", "cache.type", "copy"])
@@ -474,7 +474,7 @@ def test_should_update_state_entry_for_file_after_add(mocker, dvc, tmp_dir):
 def test_should_update_state_entry_for_directory_after_add(
     mocker, dvc, tmp_dir
 ):
-    file_md5_counter = mocker.spy(dvc_module.oid, "file_md5")
+    file_md5_counter = mocker.spy(dvc_module.objects.stage, "file_md5")
 
     tmp_dir.gen({"data/data": "foo", "data/data_sub/sub_data": "foo"})
 
@@ -521,7 +521,7 @@ class TestAddCommit(TestDvc):
 
 def test_should_collect_dir_cache_only_once(mocker, tmp_dir, dvc):
     tmp_dir.gen({"data/data": "foo"})
-    get_dir_hash_counter = mocker.spy(dvc_module.oid, "get_dir_hash")
+    get_dir_hash_counter = mocker.spy(dvc_module.objects.stage, "get_dir_hash")
     ret = main(["add", "data"])
     assert ret == 0
 

tests/func/test_data_cloud.py

@@ -231,7 +231,7 @@ def test_warn_on_outdated_stage(tmp_dir, dvc, local_remote, caplog):
 
 def test_hash_recalculation(mocker, dvc, tmp_dir, local_remote):
     tmp_dir.gen({"foo": "foo"})
-    test_get_file_hash = mocker.spy(dvc_module.oid, "get_file_hash")
+    test_get_file_hash = mocker.spy(dvc_module.objects.stage, "get_file_hash")
     ret = main(["config", "cache.type", "hardlink"])
     assert ret == 0
     ret = main(["add", "foo"])
@@ -287,7 +287,7 @@ def test_verify_hashes(
     remove("dir")
     remove(dvc.cache.local.cache_dir)
 
-    hash_spy = mocker.spy(dvc_module.oid, "get_file_hash")
+    hash_spy = mocker.spy(dvc_module.objects.stage, "get_file_hash")
 
     dvc.pull()
     assert hash_spy.call_count == 0

tests/func/test_external_repo.py

@@ -4,7 +4,8 @@ from mock import ANY, patch
 
 from dvc.external_repo import CLONES, external_repo
 from dvc.fs.local import LocalFileSystem
-from dvc.objects import save, stage
+from dvc.objects import save
+from dvc.objects.stage import stage
 from dvc.path_info import PathInfo
 from dvc.scm.git import Git
 from dvc.utils import relpath

tests/func/test_remote.py

@@ -146,7 +146,7 @@ class TestRemoteShouldHandleUppercaseRemoteName(TestDvc):
 
 
 def test_dir_hash_should_be_key_order_agnostic(tmp_dir, dvc):
-    from dvc.oid import get_hash
+    from dvc.objects.stage import get_hash
 
     tmp_dir.gen({"data": {"1": "1 content", "2": "2 content"}})
 
@@ -156,7 +156,7 @@ def test_dir_hash_should_be_key_order_agnostic(tmp_dir, dvc):
         [{"relpath": "1", "md5": "1"}, {"relpath": "2", "md5": "2"}]
     )
     with patch(
-        "dvc.oid._collect_dir", return_value=dir_info,
+        "dvc.objects.stage._collect_dir", return_value=dir_info,
     ):
         hash1 = get_hash(path_info, dvc.cache.local.fs, "md5")
 
@@ -164,7 +164,7 @@ def test_dir_hash_should_be_key_order_agnostic(tmp_dir, dvc):
         [{"md5": "1", "relpath": "1"}, {"md5": "2", "relpath": "2"}]
     )
     with patch(
-        "dvc.oid._collect_dir", return_value=dir_info,
+        "dvc.objects.stage._collect_dir", return_value=dir_info,
     ):
         hash2 = get_hash(path_info, dvc.cache.local.fs, "md5")
 

tests/unit/fs/test_dvc.py

@@ -5,7 +5,7 @@ import pytest
 
 from dvc.fs.dvc import DvcFileSystem
 from dvc.hash_info import HashInfo
-from dvc.oid import get_hash
+from dvc.objects.stage import get_hash
 from dvc.path_info import PathInfo
 
 
@@ -214,7 +214,7 @@ def test_get_hash_dir(tmp_dir, dvc, mocker):
         {"dir": {"foo": "foo", "bar": "bar", "subdir": {"data": "data"}}}
     )
     fs = DvcFileSystem(dvc)
-    get_file_hash_spy = mocker.spy(dvc_module.oid, "get_file_hash")
+    get_file_hash_spy = mocker.spy(dvc_module.objects.stage, "get_file_hash")
     assert (
         fs.info(PathInfo(tmp_dir) / "dir")["md5"]
         == "8761c4e9acad696bee718615e23e22db.dir"

tests/unit/fs/test_repo.py

@@ -6,7 +6,7 @@ import pytest
 
 from dvc.fs.repo import RepoFileSystem
 from dvc.hash_info import HashInfo
-from dvc.oid import get_hash
+from dvc.objects.stage import get_hash
 from dvc.path_info import PathInfo
 
 

tests/unit/test_external_repo.py

@@ -4,7 +4,8 @@ from unittest.mock import call
 import pytest
 
 from dvc.external_repo import external_repo
-from dvc.objects import save, stage
+from dvc.objects import save
+from dvc.objects.stage import stage
 from dvc.path_info import PathInfo
 from tests.unit.fs.test_repo import make_subrepo