index: cache target collection

dvc/repo/index.py

@@ -1,13 +1,13 @@
 from contextlib import suppress
 from typing import (
     TYPE_CHECKING,
+    Any,
     Callable,
     Dict,
     Iterable,
     Iterator,
     List,
     Optional,
-    Sequence,
     Set,
 )
 
@@ -22,7 +22,7 @@ if TYPE_CHECKING:
     from dvc.dependency import Dependency, ParamsDependency
     from dvc.fs import FileSystem
     from dvc.output import Output
-    from dvc.repo.stage import StageLoad
+    from dvc.repo.stage import StageInfo, StageLoad
     from dvc.stage import Stage
     from dvc.types import StrPath, TargetType
     from dvc_data.hashfile.hash_info import HashInfo
@@ -64,6 +64,7 @@ class Index:
         self.stage_collector: "StageLoad" = repo.stage
         if stages is not None:
             self.stages: List["Stage"] = stages
+        self._collected_targets: Dict[int, List["StageInfo"]] = {}
 
     @cached_property
     def stages(self) -> List["Stage"]:  # pylint: disable=method-hidden
@@ -113,6 +114,40 @@ class Index:
 
         return self.filter(is_stage_inside_path)
 
+    @staticmethod
+    def _hash_targets(
+        targets: Iterable[Optional[str]],
+        **kwargs: Any,
+    ) -> int:
+        return hash(
+            (
+                frozenset(targets),
+                kwargs.get("with_deps", False),
+                kwargs.get("recursive", False),
+            )
+        )
+
+    def _collect_targets(
+        self, targets: Optional["TargetType"], **kwargs: Any
+    ) -> List["StageInfo"]:
+        from itertools import chain
+
+        from dvc.repo.stage import StageInfo
+        from dvc.utils.collections import ensure_list
+
+        targets = ensure_list(targets)
+        if not targets:
+            return [StageInfo(stage) for stage in self.stages]
+        targets_hash = self._hash_targets(targets, **kwargs)
+        if targets_hash not in self._collected_targets:
+            self._collected_targets[targets_hash] = list(
+                chain.from_iterable(
+                    self.stage_collector.collect_granular(target, **kwargs)
+                    for target in targets
+                )
+            )
+        return self._collected_targets[targets_hash]
+
     @property
     def outs(self) -> Iterator["Output"]:
         for stage in self:
@@ -216,22 +251,11 @@ class Index:
         jobs: int = None,
     ) -> "ObjectContainer":
         from collections import defaultdict
-        from itertools import chain
-
-        from dvc.utils.collections import ensure_list
 
         used: "ObjectContainer" = defaultdict(set)
-        collect_targets: Sequence[Optional[str]] = (None,)
-        if targets:
-            collect_targets = ensure_list(targets)
-
-        pairs = chain.from_iterable(
-            self.stage_collector.collect_granular(
-                target, recursive=recursive, with_deps=with_deps
-            )
-            for target in collect_targets
+        pairs = self._collect_targets(
+            targets, recursive=recursive, with_deps=with_deps
         )
-
         for stage, filter_info in pairs:
             for odb, objs in stage.get_used_objs(
                 remote=remote,
@@ -247,19 +271,7 @@ class Index:
         targets: "TargetType" = None,
         recursive: bool = False,
     ) -> List["Stage"]:
-        from itertools import chain
-
-        from dvc.utils.collections import ensure_list
-
-        collect_targets: Sequence[Optional[str]] = (None,)
-        if targets:
-            collect_targets = ensure_list(targets)
-
-        pairs = chain.from_iterable(
-            self.stage_collector.collect_granular(target, recursive=recursive)
-            for target in collect_targets
-        )
-
+        pairs = self._collect_targets(targets, recursive=recursive)
         return [stage for stage, _ in pairs if stage.is_partial_import]
 
     # Following methods help us treat the collection as a set-like structure