From fd76e0991f6dd47e554b874c2efb8c280dee96c0 Mon Sep 17 00:00:00 2001
From: Saugat Pachhai <suagatchhetri@outlook.com>
Date: Fri, 6 Aug 2021 15:35:01 +0545
Subject: [PATCH] cleanup index, and add test for used_objs (#6390)

---
 dvc/repo/index.py             |  3 +--
 dvc/repo/run.py               |  3 ---
 tests/func/test_repo_index.py | 51 ++++++++++++++++++++++++++++++++++-
 3 files changed, 51 insertions(+), 6 deletions(-)

diff --git a/dvc/repo/index.py b/dvc/repo/index.py
index ca55ecd32..8989f78ec 100644
--- a/dvc/repo/index.py
+++ b/dvc/repo/index.py
@@ -110,7 +110,7 @@ class Index:
             yield from stage.outs
 
     @property
-    def decorated_outputs(self) -> Iterator["Output"]:
+    def decorated_outs(self) -> Iterator["Output"]:
         for output in self.outs:
             if output.is_decorated:
                 yield output
@@ -218,7 +218,6 @@ class Index:
         return self.remove(stage, ignore_not_existing=True)
 
     def difference(self, stages: Iterable["Stage"]) -> "Index":
-        # this does not preserve the order
         stages_set = set(self.stages) - set(stages)
         return Index(self.repo, self.fs, stages=list(stages_set))
 
diff --git a/dvc/repo/run.py b/dvc/repo/run.py
index 3f149d52d..c44bf2558 100644
--- a/dvc/repo/run.py
+++ b/dvc/repo/run.py
@@ -32,8 +32,5 @@ def run(
     else:
         stage.run(no_commit=no_commit, run_cache=run_cache)
 
-    new_index = self.index.add(stage)
-    new_index.check_graph()
-
     stage.dump(update_lock=not no_exec)
     return stage
diff --git a/tests/func/test_repo_index.py b/tests/func/test_repo_index.py
index 0224d6eb0..79cebfd0f 100644
--- a/tests/func/test_repo_index.py
+++ b/tests/func/test_repo_index.py
@@ -123,7 +123,7 @@ def test_deps_outs_getters(tmp_dir, dvc, run_copy_metrics):
 
     assert outputs_equal(index.outs, expected_outs)
     assert outputs_equal(index.deps, expected_deps)
-    assert outputs_equal(index.decorated_outputs, [metrics, plots])
+    assert outputs_equal(index.decorated_outs, [metrics, plots])
     assert outputs_equal(index.metrics, [metrics])
     assert outputs_equal(index.plots, [plots])
     assert outputs_equal(index.params, [params])
@@ -245,3 +245,52 @@ def test_skip_graph_checks(dvc, mocker):
     dvc._skip_graph_checks = False
     Index(dvc).check_graph()
     assert mock_build_graph.called
+
+
+def get_index(dvc, rev):
+    if rev:
+        brancher = dvc.brancher(revs=[rev])
+        if rev != "workspace":
+            assert next(brancher) == "workspace"
+        next(brancher)
+    return Index(dvc)
+
+
+@pytest.mark.parametrize("rev", ["workspace", "HEAD"])
+def test_used_objs(tmp_dir, scm, dvc, run_copy, rev):
+    from dvc.hash_info import HashInfo
+
+    dvc.config["core"]["autostage"] = True
+    tmp_dir.dvc_gen({"dir": {"subdir": {"file": "file"}}, "foo": "foo"})
+    run_copy("foo", "bar", name="copy-foo-bar")
+    scm.commit("commit")
+
+    index = get_index(dvc, rev)
+
+    expected_objs = [
+        HashInfo(
+            name="md5",
+            value="acbd18db4cc2f85cedef654fccc4a4d8",
+            size=3,
+            obj_name="bar",
+        ),
+        HashInfo(
+            name="md5",
+            value="8c7dd922ad47494fc02c388e12c00eac",
+            obj_name="dir/subdir/file",
+        ),
+        HashInfo(
+            name="md5",
+            value="d28c9e28591aeb7e303dc6772ffa6f6b.dir",
+            size=4,
+            nfiles=1,
+            obj_name="dir",
+        ),
+    ]
+
+    assert index.used_objs() == {None: set(expected_objs)}
+    assert index.used_objs("dir") == {None: set(expected_objs[1:])}
+    assert index.used_objs(".", recursive=True) == {None: set(expected_objs)}
+    assert index.used_objs("copy-foo-bar", with_deps=True) == {
+        None: {expected_objs[0]}
+    }