state: use granularly (#4676)

* state: add simple reentrance

* state: pass Repo instead of LocalCache

* state: use granularly

Users could still load state preemptively if they want to bulk up
the operations, but other than that tree/cache will just load it
when they need it.

tests/func/test_commit.py

@@ -24,19 +24,16 @@ def test_commit_force(tmp_dir, dvc):
     tmp_dir.gen({"dir": {"file": "text1", "file2": "text2"}})
     (stage,) = dvc.add("dir", no_commit=True)
 
-    with dvc.state:
-        assert stage.outs[0].changed_cache()
+    assert stage.outs[0].changed_cache()
 
     tmp_dir.gen("dir/file", "file content modified")
 
-    with dvc.state:
-        assert stage.outs[0].changed_cache()
+    assert stage.outs[0].changed_cache()
 
     with pytest.raises(StageCommitError):
         dvc.commit(stage.path)
 
-    with dvc.state:
-        assert stage.outs[0].changed_cache()
+    assert stage.outs[0].changed_cache()
 
     dvc.commit(stage.path, force=True)
     assert dvc.status([stage.path]) == {}
@@ -53,14 +50,12 @@ def test_commit_with_deps(tmp_dir, dvc, run_copy, run_kw):
     assert stage is not None
     assert len(stage.outs) == 1
 
-    with dvc.state:
-        assert foo_stage.outs[0].changed_cache()
-        assert stage.outs[0].changed_cache()
+    assert foo_stage.outs[0].changed_cache()
+    assert stage.outs[0].changed_cache()
 
     dvc.commit(stage.path, with_deps=True)
-    with dvc.state:
-        assert not foo_stage.outs[0].changed_cache()
-        assert not stage.outs[0].changed_cache()
+    assert not foo_stage.outs[0].changed_cache()
+    assert not stage.outs[0].changed_cache()
 
 
 def test_commit_changed_md5(tmp_dir, dvc):
@@ -125,5 +120,4 @@ def test_imported_entries_unchanged(tmp_dir, dvc, erepo_dir):
 
     stage = dvc.imp(fspath(erepo_dir), "file")
 
-    with dvc.state:
-        assert stage.changed_entries() == ([], [], None)
+    assert stage.changed_entries() == ([], [], None)

tests/func/test_data_cloud.py

@@ -55,76 +55,75 @@ def test_cloud(tmp_dir, dvc, remote):  # pylint:disable=unused-argument
     md5_dir = out_dir.hash_info.value
     info_dir = NamedCache.make(out_dir.scheme, md5_dir, name_dir)
 
-    with dvc.state:
-        # Check status
-        status = dvc.cloud.status(info, show_checksums=True)
-        expected = {md5: {"name": md5, "status": STATUS_NEW}}
-        assert status == expected
+    # Check status
+    status = dvc.cloud.status(info, show_checksums=True)
+    expected = {md5: {"name": md5, "status": STATUS_NEW}}
+    assert status == expected
 
-        status_dir = dvc.cloud.status(info_dir, show_checksums=True)
-        expected = {md5_dir: {"name": md5_dir, "status": STATUS_NEW}}
-        assert status_dir == expected
+    status_dir = dvc.cloud.status(info_dir, show_checksums=True)
+    expected = {md5_dir: {"name": md5_dir, "status": STATUS_NEW}}
+    assert status_dir == expected
 
-        # Move cache and check status
-        # See issue https://github.com/iterative/dvc/issues/4383 for details
-        backup_dir = dvc.cache.local.cache_dir + ".backup"
-        move(dvc.cache.local.cache_dir, backup_dir)
-        status = dvc.cloud.status(info, show_checksums=True)
-        expected = {md5: {"name": md5, "status": STATUS_MISSING}}
-        assert status == expected
+    # Move cache and check status
+    # See issue https://github.com/iterative/dvc/issues/4383 for details
+    backup_dir = dvc.cache.local.cache_dir + ".backup"
+    move(dvc.cache.local.cache_dir, backup_dir)
+    status = dvc.cloud.status(info, show_checksums=True)
+    expected = {md5: {"name": md5, "status": STATUS_MISSING}}
+    assert status == expected
 
-        status_dir = dvc.cloud.status(info_dir, show_checksums=True)
-        expected = {md5_dir: {"name": md5_dir, "status": STATUS_MISSING}}
-        assert status_dir == expected
+    status_dir = dvc.cloud.status(info_dir, show_checksums=True)
+    expected = {md5_dir: {"name": md5_dir, "status": STATUS_MISSING}}
+    assert status_dir == expected
 
-        # Restore original cache:
-        remove(dvc.cache.local.cache_dir)
-        move(backup_dir, dvc.cache.local.cache_dir)
+    # Restore original cache:
+    remove(dvc.cache.local.cache_dir)
+    move(backup_dir, dvc.cache.local.cache_dir)
 
-        # Push and check status
-        dvc.cloud.push(info)
-        assert os.path.exists(cache)
-        assert os.path.isfile(cache)
+    # Push and check status
+    dvc.cloud.push(info)
+    assert os.path.exists(cache)
+    assert os.path.isfile(cache)
 
-        dvc.cloud.push(info_dir)
-        assert os.path.isfile(cache_dir)
+    dvc.cloud.push(info_dir)
+    assert os.path.isfile(cache_dir)
 
-        status = dvc.cloud.status(info, show_checksums=True)
-        expected = {md5: {"name": md5, "status": STATUS_OK}}
-        assert status == expected
+    status = dvc.cloud.status(info, show_checksums=True)
+    expected = {md5: {"name": md5, "status": STATUS_OK}}
+    assert status == expected
 
-        status_dir = dvc.cloud.status(info_dir, show_checksums=True)
-        expected = {md5_dir: {"name": md5_dir, "status": STATUS_OK}}
-        assert status_dir == expected
+    status_dir = dvc.cloud.status(info_dir, show_checksums=True)
+    expected = {md5_dir: {"name": md5_dir, "status": STATUS_OK}}
+    assert status_dir == expected
 
-        # Remove and check status
-        remove(dvc.cache.local.cache_dir)
+    # Remove and check status
+    remove(dvc.cache.local.cache_dir)
 
-        status = dvc.cloud.status(info, show_checksums=True)
-        expected = {md5: {"name": md5, "status": STATUS_DELETED}}
-        assert status == expected
+    status = dvc.cloud.status(info, show_checksums=True)
+    expected = {md5: {"name": md5, "status": STATUS_DELETED}}
+    assert status == expected
 
-        status_dir = dvc.cloud.status(info_dir, show_checksums=True)
-        expected = {md5_dir: {"name": md5_dir, "status": STATUS_DELETED}}
-        assert status_dir == expected
+    status_dir = dvc.cloud.status(info_dir, show_checksums=True)
+    expected = {md5_dir: {"name": md5_dir, "status": STATUS_DELETED}}
+    assert status_dir == expected
 
-        # Pull and check status
-        dvc.cloud.pull(info)
-        assert os.path.exists(cache)
-        assert os.path.isfile(cache)
-        with open(cache) as fd:
-            assert fd.read() == "foo"
+    # Pull and check status
+    dvc.cloud.pull(info)
+    assert os.path.exists(cache)
+    assert os.path.isfile(cache)
+    with open(cache) as fd:
+        assert fd.read() == "foo"
 
-        dvc.cloud.pull(info_dir)
-        assert os.path.isfile(cache_dir)
+    dvc.cloud.pull(info_dir)
+    assert os.path.isfile(cache_dir)
 
-        status = dvc.cloud.status(info, show_checksums=True)
-        expected = {md5: {"name": md5, "status": STATUS_OK}}
-        assert status == expected
+    status = dvc.cloud.status(info, show_checksums=True)
+    expected = {md5: {"name": md5, "status": STATUS_OK}}
+    assert status == expected
 
-        status_dir = dvc.cloud.status(info_dir, show_checksums=True)
-        expected = {md5_dir: {"name": md5_dir, "status": STATUS_OK}}
-        assert status_dir == expected
+    status_dir = dvc.cloud.status(info_dir, show_checksums=True)
+    expected = {md5_dir: {"name": md5_dir, "status": STATUS_OK}}
+    assert status_dir == expected
 
 
 @pytest.mark.parametrize("remote", all_clouds, indirect=True)

tests/func/test_remote.py

@@ -149,26 +149,25 @@ def test_dir_hash_should_be_key_order_agnostic(tmp_dir, dvc):
     tmp_dir.gen({"data": {"1": "1 content", "2": "2 content"}})
 
     path_info = PathInfo("data")
-    with dvc.state:
-        with patch.object(
-            BaseTree,
-            "_collect_dir",
-            return_value=[
-                {"relpath": "1", "md5": "1"},
-                {"relpath": "2", "md5": "2"},
-            ],
-        ):
-            hash1 = dvc.cache.local.tree.get_hash(path_info)
+    with patch.object(
+        BaseTree,
+        "_collect_dir",
+        return_value=[
+            {"relpath": "1", "md5": "1"},
+            {"relpath": "2", "md5": "2"},
+        ],
+    ):
+        hash1 = dvc.cache.local.tree.get_hash(path_info)
 
-        with patch.object(
-            BaseTree,
-            "_collect_dir",
-            return_value=[
-                {"md5": "1", "relpath": "1"},
-                {"md5": "2", "relpath": "2"},
-            ],
-        ):
-            hash2 = dvc.cache.local.tree.get_hash(path_info)
+    with patch.object(
+        BaseTree,
+        "_collect_dir",
+        return_value=[
+            {"md5": "1", "relpath": "1"},
+            {"md5": "2", "relpath": "2"},
+        ],
+    ):
+        hash2 = dvc.cache.local.tree.get_hash(path_info)
 
     assert hash1 == hash2
 

tests/func/test_run_cache.py

@@ -46,8 +46,7 @@ def test_save(tmp_dir, dvc, run_copy):
     tmp_dir.gen("foo", "foo")
     stage = run_copy("foo", "bar", name="copy-foo-bar")
     assert _recurse_count_files(run_cache_dir) == 1
-    with dvc.state:
-        assert dvc.stage_cache._load(stage)
+    assert dvc.stage_cache._load(stage)
 
 
 def test_do_not_save_on_no_exec_and_dry(tmp_dir, dvc, run_copy):
@@ -58,14 +57,12 @@ def test_do_not_save_on_no_exec_and_dry(tmp_dir, dvc, run_copy):
     stage = run_copy("foo", "bar", name="copy-foo-bar", no_exec=True)
 
     assert _recurse_count_files(run_cache_dir) == 0
-    with dvc.state:
-        assert not dvc.stage_cache._load(stage)
+    assert not dvc.stage_cache._load(stage)
 
     (stage,) = dvc.reproduce("copy-foo-bar", dry=True)
 
     assert _recurse_count_files(run_cache_dir) == 0
-    with dvc.state:
-        assert not dvc.stage_cache._load(stage)
+    assert not dvc.stage_cache._load(stage)
 
 
 def test_uncached_outs_are_cached(tmp_dir, dvc, run_copy):
@@ -76,8 +73,7 @@ def test_uncached_outs_are_cached(tmp_dir, dvc, run_copy):
         outs_no_cache=["bar"],
         name="copy-foo-bar",
     )
-    with dvc.state:
-        stage.outs[0].hash_info = stage.outs[0].get_hash()
+    stage.outs[0].hash_info = stage.outs[0].get_hash()
     assert os.path.exists(relpath(stage.outs[0].cache_path))
 
 

tests/func/test_state.py

@@ -13,7 +13,7 @@ def test_state(tmp_dir, dvc):
     path_info = PathInfo(path)
     md5 = file_md5(path)[0]
 
-    state = State(dvc.cache.local)
+    state = State(dvc)
 
     with state:
         state.save(path_info, md5)
@@ -57,7 +57,7 @@ def mock_get_inode(inode):
 def test_get_state_record_for_inode(get_inode_mock, tmp_dir, dvc):
     tmp_dir.gen("foo", "foo content")
 
-    state = State(dvc.cache.local)
+    state = State(dvc)
     inode = state.MAX_INT + 2
     assert inode != state._to_sqlite(inode)
 

tests/func/test_tree.py

@@ -194,9 +194,8 @@ def test_repotree_walk_fetch(tmp_dir, dvc, scm, local_remote):
     remove(tmp_dir / "dir")
 
     tree = RepoTree(dvc, fetch=True)
-    with dvc.state:
-        for _, _, _ in tree.walk("dir"):
-            pass
+    for _, _, _ in tree.walk("dir"):
+        pass
 
     assert os.path.exists(out.cache_path)
     for entry in out.dir_cache:
@@ -223,12 +222,10 @@ def test_repotree_cache_save(tmp_dir, dvc, scm, erepo_dir, local_cloud):
         for path in ("dir/bar", "dir/subdir/foo")
     ]
 
-    with erepo_dir.dvc.state:
-        cache = dvc.cache.local
-        with cache.tree.state:
-            path_info = PathInfo(erepo_dir / "dir")
-            hash_info = cache.tree.get_hash(path_info)
-            cache.save(path_info, tree, hash_info)
+    cache = dvc.cache.local
+    path_info = PathInfo(erepo_dir / "dir")
+    hash_info = cache.tree.get_hash(path_info)
+    cache.save(path_info, tree, hash_info)
 
     for hash_ in expected:
         assert os.path.exists(cache.tree.hash_to_path_info(hash_))

tests/unit/dependency/test_params.py

@@ -208,5 +208,4 @@ def test_params_with_false_values(tmp_dir, dvc, param_value):
 
     dep.fill_values(load_yaml(DEFAULT_PARAMS_FILE))
 
-    with dvc.state:
-        assert dep.status() == {}
+    assert dep.status() == {}

tests/unit/repo/test_repo.py

@@ -50,12 +50,11 @@ def test_used_cache(tmp_dir, dvc, path):
         ),
     )
 
-    with dvc.state:
-        used_cache = dvc.used_cache([path])
-        assert (
-            used_cache._items == expected._items
-            and used_cache.external == expected.external
-        )
+    used_cache = dvc.used_cache([path])
+    assert (
+        used_cache._items == expected._items
+        and used_cache.external == expected.external
+    )
 
 
 def test_locked(mocker):

tests/unit/tree/test_dvc.py

@@ -23,9 +23,8 @@ def test_open(tmp_dir, dvc):
     (tmp_dir / "foo").unlink()
 
     tree = DvcTree(dvc)
-    with dvc.state:
-        with tree.open("foo", "r") as fobj:
-            assert fobj.read() == "foo"
+    with tree.open("foo", "r") as fobj:
+        assert fobj.read() == "foo"
 
 
 def test_open_dirty_hash(tmp_dir, dvc):
@@ -169,11 +168,10 @@ def test_walk_dir(tmp_dir, dvc, fetch, expected):
 
     expected = [str(tmp_dir / path) for path in expected]
 
-    with dvc.state:
-        actual = []
-        for root, dirs, files in tree.walk("dir"):
-            for entry in dirs + files:
-                actual.append(os.path.join(root, entry))
+    actual = []
+    for root, dirs, files in tree.walk("dir"):
+        for entry in dirs + files:
+            actual.append(os.path.join(root, entry))
 
     assert set(actual) == set(expected)
     assert len(actual) == len(expected)
@@ -235,13 +233,12 @@ def test_get_hash_granular(tmp_dir, dvc):
     )
     tree = DvcTree(dvc, fetch=True)
     subdir = PathInfo(tmp_dir) / "dir" / "subdir"
-    with dvc.state:
-        assert tree.get_hash(subdir) == HashInfo(
-            "md5", "af314506f1622d107e0ed3f14ec1a3b5.dir",
-        )
-        assert tree.get_hash(subdir / "data") == HashInfo(
-            "md5", "8d777f385d3dfec8815d20f7496026dc",
-        )
+    assert tree.get_hash(subdir) == HashInfo(
+        "md5", "af314506f1622d107e0ed3f14ec1a3b5.dir",
+    )
+    assert tree.get_hash(subdir / "data") == HashInfo(
+        "md5", "8d777f385d3dfec8815d20f7496026dc",
+    )
 
 
 def test_get_hash_dirty_file(tmp_dir, dvc):
@@ -249,8 +246,7 @@ def test_get_hash_dirty_file(tmp_dir, dvc):
     (tmp_dir / "file").write_text("something")
 
     tree = DvcTree(dvc)
-    with dvc.state:
-        actual = tree.get_hash(PathInfo(tmp_dir) / "file")
+    actual = tree.get_hash(PathInfo(tmp_dir) / "file")
     expected = HashInfo("md5", "8c7dd922ad47494fc02c388e12c00eac")
     assert actual == expected
 
@@ -260,8 +256,7 @@ def test_get_hash_dirty_dir(tmp_dir, dvc):
     (tmp_dir / "dir" / "baz").write_text("baz")
 
     tree = DvcTree(dvc)
-    with dvc.state:
-        actual = tree.get_hash(PathInfo(tmp_dir) / "dir")
+    actual = tree.get_hash(PathInfo(tmp_dir) / "dir")
     expected = HashInfo("md5", "5ea40360f5b4ec688df672a4db9c17d1.dir")
 
     assert actual == expected

tests/unit/tree/test_repo.py

@@ -24,9 +24,8 @@ def test_open(tmp_dir, dvc):
     (tmp_dir / "foo").unlink()
 
     tree = RepoTree(dvc)
-    with dvc.state:
-        with tree.open("foo", "r") as fobj:
-            assert fobj.read() == "foo"
+    with tree.open("foo", "r") as fobj:
+        assert fobj.read() == "foo"
 
 
 def test_open_dirty_hash(tmp_dir, dvc):
@@ -494,19 +493,17 @@ def test_get_hash_cached_dir(tmp_dir, dvc, mocker):
     tree = RepoTree(dvc)
     get_file_hash_spy = mocker.spy(tree, "get_file_hash")
     dvc_tree_spy = mocker.spy(tree._dvctrees[dvc.root_dir], "get_dir_hash")
-    with dvc.state:
-        assert tree.get_hash(PathInfo(tmp_dir) / "dir") == HashInfo(
-            "md5", "8761c4e9acad696bee718615e23e22db.dir",
-        )
+    assert tree.get_hash(PathInfo(tmp_dir) / "dir") == HashInfo(
+        "md5", "8761c4e9acad696bee718615e23e22db.dir",
+    )
     assert get_file_hash_spy.called
     assert not dvc_tree_spy.called
     get_file_hash_spy.reset_mock()
 
     shutil.rmtree(tmp_dir / "dir")
-    with dvc.state:
-        assert tree.get_hash(PathInfo(tmp_dir) / "dir") == HashInfo(
-            "md5", "8761c4e9acad696bee718615e23e22db.dir",
-        )
+    assert tree.get_hash(PathInfo(tmp_dir) / "dir") == HashInfo(
+        "md5", "8761c4e9acad696bee718615e23e22db.dir",
+    )
     assert not get_file_hash_spy.called
     assert dvc_tree_spy.called
 
@@ -518,13 +515,12 @@ def test_get_hash_cached_granular(tmp_dir, dvc, mocker):
     tree = RepoTree(dvc, fetch=True)
     dvc_tree_spy = mocker.spy(tree._dvctrees[dvc.root_dir], "get_file_hash")
     subdir = PathInfo(tmp_dir) / "dir" / "subdir"
-    with dvc.state:
-        assert tree.get_hash(subdir) == HashInfo(
-            "md5", "af314506f1622d107e0ed3f14ec1a3b5.dir",
-        )
-        assert tree.get_hash(subdir / "data") == HashInfo(
-            "md5", "8d777f385d3dfec8815d20f7496026dc",
-        )
+    assert tree.get_hash(subdir) == HashInfo(
+        "md5", "af314506f1622d107e0ed3f14ec1a3b5.dir",
+    )
+    assert tree.get_hash(subdir / "data") == HashInfo(
+        "md5", "8d777f385d3dfec8815d20f7496026dc",
+    )
     assert dvc_tree_spy.called
 
 
@@ -541,8 +537,7 @@ def test_get_hash_mixed_dir(tmp_dir, scm, dvc):
     tmp_dir.scm.commit("add dir")
 
     tree = RepoTree(dvc)
-    with dvc.state:
-        actual = tree.get_hash(PathInfo(tmp_dir) / "dir")
+    actual = tree.get_hash(PathInfo(tmp_dir) / "dir")
     expected = HashInfo("md5", "e1d9e8eae5374860ae025ec84cfd85c7.dir")
     assert actual == expected
 
@@ -552,8 +547,7 @@ def test_get_hash_dirty_file(tmp_dir, dvc):
     (tmp_dir / "file").write_text("something")
 
     tree = RepoTree(dvc)
-    with dvc.state:
-        actual = tree.get_hash(PathInfo(tmp_dir) / "file")
+    actual = tree.get_hash(PathInfo(tmp_dir) / "file")
     expected = HashInfo("md5", "8c7dd922ad47494fc02c388e12c00eac")
     assert actual == expected
 
@@ -563,8 +557,7 @@ def test_get_hash_dirty_dir(tmp_dir, dvc):
     (tmp_dir / "dir" / "baz").write_text("baz")
 
     tree = RepoTree(dvc)
-    with dvc.state:
-        actual = tree.get_hash(PathInfo(tmp_dir) / "dir")
+    actual = tree.get_hash(PathInfo(tmp_dir) / "dir")
     expected = HashInfo("md5", "ba75a2162ca9c29acecb7957105a0bc2.dir")
     assert actual == expected
     assert len(actual.dir_info) == 3