index/stage: move repo collection logic to index

dvc/dvcfile.py

@@ -1,16 +1,7 @@
 import contextlib
 import logging
 import os
-from typing import (
+from typing import TYPE_CHECKING, Any, Callable, List, Tuple, TypeVar, Union
-    TYPE_CHECKING,
-    Any,
-    Callable,
-    Dict,
-    List,
-    Tuple,
-    TypeVar,
-    Union,
-)
 
 from funcy import cached_property
 
@@ -192,7 +183,7 @@ class SingleStageFile(FileMixin):
     from dvc.stage.loader import SingleStageLoader as LOADER
 
     metrics: List[str] = []
-    plots: Dict[str, Any] = {}
+    plots: Any = {}
     params: List[str] = []
 
     @property
@@ -242,6 +233,7 @@ class PipelineFile(FileMixin):
         self.__dict__.pop("contents", None)
         self.__dict__.pop("lockfile_contents", None)
         self.__dict__.pop("resolver", None)
+        self.__dict__.pop("stages", None)
 
     def dump(
         self, stage, update_pipeline=True, update_lock=True, **kwargs
@@ -313,7 +305,7 @@ class PipelineFile(FileMixin):
         wdir = self.repo.fs.path.parent(self.path)
         return DataResolver(self.repo, wdir, self.contents)
 
-    @property
+    @cached_property
     def stages(self):
         return self.LOADER(self, self.contents, self.lockfile_contents)
 

dvc/repo/__init__.py

@@ -239,7 +239,14 @@ class Repo:
     def index(self):
         from dvc.repo.index import Index
 
-        return Index(self)
+        return Index.from_repo(self)
+
+    def ensure_graph_correctness_with(self, stages=None, callback=None):
+        if not getattr(self, "_skip_graph_checks", False):
+            new = self.index.update(stages)
+            if callable(callback):
+                callback()
+            new.check_graph()
 
     @staticmethod
     def open(url, *args, **kwargs):

dvc/repo/add.py

@@ -180,9 +180,9 @@ def add(  # noqa: C901
     with translate_graph_error(stages), ui.status(msg) as status:
         # remove existing stages that are to-be replaced with these
         # new stages for the graph checks.
-        new_index = repo.index.update(stages)
+        repo.ensure_graph_correctness_with(
-        status.update("Checking graph")
+            stages=stages, callback=lambda: status.update("Checking graph")
-        new_index.check_graph()
+        )
 
     odb = None
     if to_remote:

dvc/repo/imp_url.py

@@ -78,8 +78,7 @@ def imp_url(  # noqa: C901
     dvcfile.remove()
 
     try:
-        new_index = self.index.add(stage)
+        self.ensure_graph_correctness_with(stages={stage})
-        new_index.check_graph()
     except OutputDuplicationError as exc:
         raise OutputDuplicationError(  # noqa: B904
             exc.output, set(exc.stages) - {stage}

dvc/repo/index.py

@@ -1,3 +1,6 @@
+import logging
+import time
+from functools import partial
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -14,16 +17,18 @@ from typing import (
 )
 
 from funcy import cached_property
+from funcy.debug import format_time
+
+from dvc.fs import LocalFileSystem
 
 if TYPE_CHECKING:
     from networkx import DiGraph
     from pygtrie import Trie
 
     from dvc.dependency import Dependency, ParamsDependency
-    from dvc.fs import FileSystem
     from dvc.output import Output
     from dvc.repo import Repo
-    from dvc.repo.stage import StageInfo, StageLoad
+    from dvc.repo.stage import StageInfo
     from dvc.stage import Stage
     from dvc.types import TargetType
     from dvc_data.hashfile.hash_info import HashInfo
@@ -31,174 +36,180 @@ if TYPE_CHECKING:
     from dvc_objects.db import ObjectDB
 
 
+logger = logging.getLogger(__name__)
 ObjectContainer = Dict[Optional["ObjectDB"], Set["HashInfo"]]
 
 
+def log_walk(seq):
+    for root, dirs, files in seq:
+        start = time.perf_counter()
+        yield root, dirs, files
+        duration = format_time(time.perf_counter() - start)
+        logger.trace("%s in collecting stages from %s", duration, root)
+
+
+def collect_files(
+    repo: "Repo", onerror: Callable[[str, Exception], None] = None
+):
+    """Collects all of the stages present in the DVC repo.
+
+    Args:
+        onerror (optional): callable that will be called with two args:
+            the filepath whose collection failed and the exc instance.
+            It can report the error to continue with the collection
+            (and, skip failed ones), or raise the exception to abort
+            the collection.
+    """
+    from dvc.dvcfile import is_valid_filename
+    from dvc.exceptions import DvcException
+    from dvc.utils import relpath
+
+    scm = repo.scm
+    fs = repo.fs
+    sep = fs.sep
+    outs: Set[str] = set()
+
+    is_local_fs = isinstance(fs, LocalFileSystem)
+
+    def is_ignored(path):
+        # apply only for the local fs
+        return is_local_fs and scm.is_ignored(path)
+
+    def is_dvcfile_and_not_ignored(root, file):
+        return is_valid_filename(file) and not is_ignored(f"{root}{sep}{file}")
+
+    def is_out_or_ignored(root, directory):
+        dir_path = f"{root}{sep}{directory}"
+        # trailing slash needed to check if a directory is gitignored
+        return dir_path in outs or is_ignored(f"{dir_path}{sep}")
+
+    walk_iter = repo.dvcignore.walk(fs, repo.root_dir, followlinks=False)
+    if logger.isEnabledFor(logging.TRACE):  # type: ignore[attr-defined]
+        walk_iter = log_walk(walk_iter)
+
+    for root, dirs, files in walk_iter:
+        dvcfile_filter = partial(is_dvcfile_and_not_ignored, root)
+        for file in filter(dvcfile_filter, files):
+            file_path = fs.path.join(root, file)
+            try:
+                index = Index.from_file(repo, file_path)
+            except DvcException as exc:
+                if onerror:
+                    onerror(relpath(file_path), exc)
+                    continue
+                raise
+
+            outs.update(
+                out.fspath
+                for stage in index.stages
+                for out in stage.outs
+                if out.protocol == "local"
+            )
+            yield file_path, index
+        dirs[:] = [d for d in dirs if not is_out_or_ignored(root, d)]
+
+
 class Index:
     def __init__(
         self,
-        repo: "Repo",  # pylint: disable=redefined-outer-name
+        repo: "Repo",
-        fs: "FileSystem" = None,
         stages: List["Stage"] = None,
+        metrics: Dict[str, List[str]] = None,
+        plots: Dict[str, List[str]] = None,
+        params: Dict[str, Any] = None,
     ) -> None:
-        """Index is an immutable collection of stages.
+        self.repo = repo
-
+        self.stages = stages or []
-        Generally, Index is a complete collection of stages at a point in time.
+        self._metrics = metrics or {}
-        With "a point in time", it means it is collected from the user's
+        self._plots = plots or {}
-        workspace or a git revision.
+        self._params = params or {}
-        And, since Index is immutable, the collection is frozen in time.
-
-        Index provides multiple ways to view this collection:
-
-            stages - provides direct access to this collection
-            outputs - provides direct access to the outputs
-            objects - provides direct access to the objects
-            graph -
-            ... and many more.
-
-        Index also provides ways to slice and dice this collection.
-        Some `views` might not make sense when sliced (eg: pipelines/graph).
-        """
-
-        self.repo: "Repo" = repo
-        self.fs: "FileSystem" = fs or repo.fs
-        self.stage_collector: "StageLoad" = repo.stage
-        if stages is not None:
-            self.stages: List["Stage"] = stages
         self._collected_targets: Dict[int, List["StageInfo"]] = {}
-        self._metrics: Dict[str, List[str]] = {}
-        self._plots: Dict[str, Any] = {}
-        self._params: Dict[str, List[str]] = {}
-
-    @cached_property
-    def stages(self) -> List["Stage"]:  # pylint: disable=method-hidden
-        # note that ideally we should be keeping this in a set as it is unique,
-        # hashable and has no concept of orderliness on its own. But we depend
-        # on this to be somewhat ordered for status/metrics/plots, etc.
-        return self._collect()
-
-    @cached_property
-    def _top_metrics(self):
-        self._collect()
-        return self._metrics
-
-    @cached_property
-    def _top_plots(self):
-        self._collect()
-        return self._plots
-
-    @cached_property
-    def _plot_sources(self):
-        from dvc.repo.plots import _collect_pipeline_files
-
-        sources = []
-        for data in _collect_pipeline_files(self.repo, [], {}).values():
-            for plot_id, props in data.get("data", {}).items():
-                if isinstance(props.get("y"), dict):
-                    sources.extend(props["y"])
-                    if isinstance(props.get("x"), dict):
-                        sources.extend(props["x"])
-                else:
-                    sources.append(plot_id)
-        return sources
-
-    @cached_property
-    def _top_params(self):
-        self._collect()
-        return self._params
-
-    def _collect(self):
-        if "stages" in self.__dict__:
-            return self.stages
-
-        onerror = self.repo.stage_collection_error_handler
-
-        # pylint: disable=protected-access
-        (
-            stages,
-            metrics,
-            plots,
-            params,
-        ) = self.stage_collector._collect_all_from_repo(onerror=onerror)
-        self.stages = stages
-        self._metrics = metrics
-        self._plots = plots
-        self._params = params
-        return stages
 
     def __repr__(self) -> str:
-        from dvc.fs import LocalFileSystem
-
         rev = "workspace"
-        if not isinstance(self.fs, LocalFileSystem):
+        if not isinstance(self.repo.fs, LocalFileSystem):
             rev = self.repo.get_rev()[:7]
         return f"Index({self.repo}, fs@{rev})"
 
-    @staticmethod
+    @classmethod
-    def _hash_targets(
+    def from_repo(
-        targets: Iterable[Optional[str]],
+        cls, repo: "Repo", onerror: Callable[[str, Exception], None] = None
-        **kwargs: Any,
+    ) -> "Index":
-    ) -> int:
+        stages = []
-        return hash(
+        metrics = {}
-            (
+        plots = {}
-                frozenset(targets),
+        params = {}
-                kwargs.get("with_deps", False),
+
-                kwargs.get("recursive", False),
+        onerror = onerror or repo.stage_collection_error_handler
-            )
+        for _, idx in collect_files(repo, onerror=onerror):
+            # pylint: disable=protected-access
+            stages.extend(idx.stages)
+            metrics.update(idx._metrics)
+            plots.update(idx._plots)
+            params.update(idx._params)
+        return cls(
+            repo,
+            stages=stages,
+            metrics=metrics,
+            plots=plots,
+            params=params,
         )
 
-    def collect_targets(
+    @classmethod
-        self, targets: Optional["TargetType"], **kwargs: Any
+    def from_file(cls, repo: "Repo", path: str) -> "Index":
-    ) -> List["StageInfo"]:
+        from dvc.dvcfile import make_dvcfile
-        from itertools import chain
 
-        from dvc.repo.stage import StageInfo
+        dvcfile = make_dvcfile(repo, path)
-        from dvc.utils.collections import ensure_list
+        return cls(
+            repo,
+            stages=list(dvcfile.stages.values()),
+            metrics={path: dvcfile.metrics} if dvcfile.metrics else {},
+            plots={path: dvcfile.plots} if dvcfile.plots else {},
+            params={path: dvcfile.params} if dvcfile.params else {},
+        )
 
-        targets = ensure_list(targets)
+    def update(self, stages: Iterable["Stage"]) -> "Index":
-        if not targets:
+        stages = set(stages)
-            return [StageInfo(stage) for stage in self.stages]
+        # we remove existing stages with same hashes at first
-        targets_hash = self._hash_targets(targets, **kwargs)
+        # and then re-add the new ones later.
-        if targets_hash not in self._collected_targets:
+        stages_set = (set(self.stages) - stages) | stages
-            self._collected_targets[targets_hash] = list(
+        return self.__class__(
-                chain.from_iterable(
+            self.repo,
-                    self.stage_collector.collect_granular(target, **kwargs)
+            stages=list(stages_set),
-                    for target in targets
+            metrics=self._metrics,
-                )
+            plots=self._plots,
-            )
+            params=self._params,
-        return self._collected_targets[targets_hash]
+        )
 
-    def targets_view(
+    @cached_property
-        self,
+    def outs_trie(self) -> "Trie":
-        targets: Optional["TargetType"],
+        from dvc.repo.trie import build_outs_trie
-        stage_filter: Optional[Callable[["Stage"], bool]] = None,
-        outs_filter: Optional[Callable[["Output"], bool]] = None,
-        **kwargs: Any,
-    ) -> "IndexView":
-        """Return read-only view of index for the specified targets.
 
-        Args:
+        return build_outs_trie(self.stages)
-            targets: Targets to collect
-            stage_filter: Optional stage filter to be applied after collecting
-                targets.
-            outs_filter: Optional output filter to be applied after collecting
-                targets.
 
-        Additional kwargs will be passed into the stage collector.
+    @cached_property
+    def outs_graph(self) -> "DiGraph":
+        from dvc.repo.graph import build_outs_graph
 
-        Note:
+        return build_outs_graph(self.graph, self.outs_trie)
-            If both stage_filter and outs_filter are provided, stage_filter
+
-            will be applied first, and the resulting view will only contain
+    @cached_property
-            outputs from stages that matched stage_filter. Outputs from stages
+    def graph(self) -> "DiGraph":
-            that did not match will be excluded from the view (whether or not
+        from dvc.repo.graph import build_graph
-            the output would have matched outs_filter).
+
-        """
+        return build_graph(self.stages, self.outs_trie)
-        stage_infos = [
+
-            stage_info
+    def check_graph(self) -> None:
-            for stage_info in self.collect_targets(targets, **kwargs)
+        if not getattr(self.repo, "_skip_graph_checks", False):
-            if not stage_filter or stage_filter(stage_info.stage)
+            self.graph  # pylint: disable=pointless-statement
-        ]
+
-        return IndexView(self, stage_infos, outs_filter=outs_filter)
+    @property
+    def params(self) -> Iterator["ParamsDependency"]:
+        from dvc.dependency import ParamsDependency
+
+        for dep in self.deps:
+            if isinstance(dep, ParamsDependency):
+                yield dep
 
     @property
     def outs(self) -> Iterator["Output"]:
@@ -228,31 +239,20 @@ class Index:
         for stage in self.stages:
             yield from stage.deps
 
-    @property
-    def params(self) -> Iterator["ParamsDependency"]:
-        from dvc.dependency import ParamsDependency
-
-        for dep in self.deps:
-            if isinstance(dep, ParamsDependency):
-                yield dep
-
     @cached_property
-    def outs_trie(self) -> "Trie":
+    def _plot_sources(self) -> List[str]:
-        from dvc.repo.trie import build_outs_trie
+        from dvc.repo.plots import _collect_pipeline_files
 
-        return build_outs_trie(self.stages)
+        sources: List[str] = []
-
+        for data in _collect_pipeline_files(self.repo, [], {}).values():
-    @cached_property
+            for plot_id, props in data.get("data", {}).items():
-    def graph(self) -> "DiGraph":
+                if isinstance(props.get("y"), dict):
-        from dvc.repo.graph import build_graph
+                    sources.extend(props["y"])
-
+                    if isinstance(props.get("x"), dict):
-        return build_graph(self.stages, self.outs_trie)
+                        sources.extend(props["x"])
-
+                else:
-    @cached_property
+                    sources.append(plot_id)
-    def outs_graph(self) -> "DiGraph":
+        return sources
-        from dvc.repo.graph import build_outs_graph
-
-        return build_outs_graph(self.graph, self.outs_trie)
 
     @cached_property
     def data(self) -> "Dict[str, DataIndex]":
@@ -278,6 +278,40 @@ class Index:
 
         return dict(by_workspace)
 
+    @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.repo.stage.collect_granular(target, **kwargs)
+                    for target in targets
+                )
+            )
+        return self._collected_targets[targets_hash]
+
     def used_objs(
         self,
         targets: "TargetType" = None,
@@ -305,23 +339,34 @@ class Index:
                 used[odb].update(objs)
         return used
 
-    # Following methods help us treat the collection as a set-like structure
+    def targets_view(
-    # and provides faux-immutability.
+        self,
-    # These methods do not preserve stages order.
+        targets: Optional["TargetType"],
-
+        stage_filter: Optional[Callable[["Stage"], bool]] = None,
-    def update(self, stages: Iterable["Stage"]) -> "Index":
+        outs_filter: Optional[Callable[["Output"], bool]] = None,
-        new_stages = set(stages)
+        **kwargs: Any,
-        # we remove existing stages with same hashes at first
+    ) -> "IndexView":
-        # and then re-add the new ones later.
+        """Return read-only view of index for the specified targets.
-        stages_set = (set(self.stages) - new_stages) | new_stages
+        Args:
-        return Index(self.repo, self.fs, stages=list(stages_set))
+            targets: Targets to collect
-
+            stage_filter: Optional stage filter to be applied after collecting
-    def add(self, stage: "Stage") -> "Index":
+                targets.
-        return self.update([stage])
+            outs_filter: Optional output filter to be applied after collecting
-
+                targets.
-    def check_graph(self) -> None:
+        Additional kwargs will be passed into the stage collector.
-        if not getattr(self.repo, "_skip_graph_checks", False):
+        Note:
-            self.graph  # pylint: disable=pointless-statement
+            If both stage_filter and outs_filter are provided, stage_filter
+            will be applied first, and the resulting view will only contain
+            outputs from stages that matched stage_filter. Outputs from stages
+            that did not match will be excluded from the view (whether or not
+            the output would have matched outs_filter).
+        """
+        stage_infos = [
+            stage_info
+            for stage_info in self.collect_targets(targets, **kwargs)
+            if not stage_filter or stage_filter(stage_info.stage)
+        ]
+        return IndexView(self, stage_infos, outs_filter=outs_filter)
 
 
 class _DataPrefixes(NamedTuple):
@@ -342,13 +387,9 @@ class IndexView:
         self._stage_infos = stage_infos
         # NOTE: stage_infos might have the same stage multiple times but with
         # different filter_info
-        self._stages = list({stage for stage, _ in stage_infos})
+        self.stages = list({stage for stage, _ in stage_infos})
         self._outs_filter = outs_filter
 
-    @property
-    def stages(self):
-        return self._stages
-
     @property
     def deps(self) -> Iterator["Dependency"]:
         for stage in self.stages:
@@ -385,8 +426,6 @@ class IndexView:
 
     @cached_property
     def data(self) -> Dict[str, Union["DataIndex", "DataIndexView"]]:
-        from functools import partial
-
         from dvc_data.index import DataIndex, view
 
         def key_filter(workspace: str, key: "DataIndexKey"):
@@ -401,7 +440,7 @@ class IndexView:
 
         data: Dict[str, Union["DataIndex", "DataIndexView"]] = {}
         for workspace, data_index in self._index.data.items():
-            if self._stages:
+            if self.stages:
                 data[workspace] = view(
                     data_index, partial(key_filter, workspace)
                 )

dvc/repo/metrics/show.py

@@ -34,7 +34,7 @@ def _to_fs_paths(metrics: List[Output]) -> StrPaths:
 
 
 def _collect_top_level_metrics(repo):
-    top_metrics = repo.index._top_metrics  # pylint: disable=protected-access
+    top_metrics = repo.index._metrics  # pylint: disable=protected-access
     for dvcfile, metrics in top_metrics.items():
         wdir = repo.fs.path.relpath(
             repo.fs.path.parent(dvcfile), repo.root_dir

dvc/repo/params/show.py

@@ -41,7 +41,7 @@ def _is_params(dep: "Output"):
 
 
 def _collect_top_level_params(repo):
-    top_params = repo.index._top_params  # pylint: disable=protected-access
+    top_params = repo.index._params  # pylint: disable=protected-access
     for dvcfile, params in top_params.items():
         wdir = repo.fs.path.relpath(
             repo.fs.path.parent(dvcfile), repo.root_dir

dvc/repo/plots/__init__.py

@@ -453,7 +453,7 @@ def _resolve_definitions(
 
 def _collect_pipeline_files(repo, targets: List[str], props, onerror=None):
     result: Dict[str, Dict] = {}
-    top_plots = repo.index._top_plots  # pylint: disable=protected-access
+    top_plots = repo.index._plots  # pylint: disable=protected-access
     for dvcfile, plots_def in top_plots.items():
         dvcfile_path = _relpath(repo.dvcfs, dvcfile)
         dvcfile_defs_dict: Dict[str, Union[Dict, None]] = {}

dvc/repo/stage.py

@@ -1,30 +1,14 @@
 import fnmatch
 import logging
-import time
 import typing
 from contextlib import suppress
-from functools import partial, wraps
+from functools import wraps
-from typing import (
+from typing import Iterable, List, NamedTuple, Optional, Set, Tuple, Union
-    Callable,
-    Iterable,
-    List,
-    NamedTuple,
-    Optional,
-    Set,
-    Tuple,
-    Union,
-)
 
-from funcy.debug import format_time
+from dvc.exceptions import NoOutputOrStageError, OutputNotFoundError
-
-from dvc.exceptions import (
-    DvcException,
-    NoOutputOrStageError,
-    OutputNotFoundError,
-)
 from dvc.repo import lock_repo
 from dvc.ui import ui
-from dvc.utils import as_posix, parse_target, relpath
+from dvc.utils import as_posix, parse_target
 
 logger = logging.getLogger(__name__)
 
@@ -39,14 +23,6 @@ if typing.TYPE_CHECKING:
 PIPELINE_FILE = "dvc.yaml"
 
 
-def log_walk(seq):
-    for root, dirs, files in seq:
-        start = time.perf_counter()
-        yield root, dirs, files
-        duration = format_time(time.perf_counter() - start)
-        logger.trace("%s in collecting stages from %s", duration, root)
-
-
 class StageInfo(NamedTuple):
     stage: "Stage"
     filter_info: Optional[str] = None
@@ -75,7 +51,9 @@ def _maybe_collect_from_dvc_yaml(
     if loader.fs.exists(PIPELINE_FILE):
         with suppress(StageNotFound):
             stages = loader.load_all(PIPELINE_FILE, target, **load_kwargs)
-    return _collect_with_deps(stages, loader.graph) if with_deps else stages
+    if with_deps:
+        return _collect_with_deps(stages, loader.repo.index.graph)
+    return stages
 
 
 def _collect_specific_target(
@@ -117,9 +95,9 @@ def locked(f):
 
 
 class StageLoad:
-    def __init__(self, repo: "Repo", fs=None) -> None:
+    def __init__(self, repo: "Repo") -> None:
         self.repo: "Repo" = repo
-        self._fs = fs
+        self.fs = repo.fs
 
     @locked
     def add(
@@ -207,8 +185,7 @@ class StageLoad:
 
                 check_stage_exists(self.repo, stage, stage.path)
 
-            new_index = self.repo.index.add(stage)
+            self.repo.ensure_graph_correctness_with(stages={stage})
-            new_index.check_graph()
 
         restore_fields(stage)
         return stage
@@ -323,16 +300,6 @@ class StageLoad:
         """Load stages from `path`, filtered with `expr` provided."""
         return self.load_all(path, expr, glob=True)
 
-    @property
-    def fs(self):
-        if self._fs:
-            return self._fs
-        return self.repo.fs
-
-    @property
-    def graph(self) -> "DiGraph":
-        return self.repo.index.graph
-
     def collect(
         self,
         target: str = None,
@@ -375,13 +342,13 @@ class StageLoad:
             from dvc.repo.graph import collect_inside_path
 
             path = self.fs.path.abspath(target)
-            return collect_inside_path(path, graph or self.graph)
+            return collect_inside_path(path, graph or self.repo.index.graph)
 
         stages = self.from_target(target, glob=glob)
         if not with_deps:
             return stages
 
-        return _collect_with_deps(stages, graph or self.graph)
+        return _collect_with_deps(stages, graph or self.repo.index.graph)
 
     def collect_granular(
         self,
@@ -449,153 +416,3 @@ class StageLoad:
                 raise NoOutputOrStageError(target, exc.file) from exc
 
         return [StageInfo(stage) for stage in stages]
-
-    def _collect_repo(self, onerror: Callable[[str, Exception], None] = None):
-        """Collects all of the stages present in the DVC repo.
-
-        Args:
-            onerror (optional): callable that will be called with two args:
-                the filepath whose collection failed and the exc instance.
-                It can report the error to continue with the collection
-                (and, skip failed ones), or raise the exception to abort
-                the collection.
-        """
-        from dvc.dvcfile import is_valid_filename
-        from dvc.fs import LocalFileSystem
-
-        scm = self.repo.scm
-        sep = self.fs.sep
-        outs: Set[str] = set()
-
-        is_local_fs = isinstance(self.fs, LocalFileSystem)
-
-        def is_ignored(path):
-            # apply only for the local fs
-            return is_local_fs and scm.is_ignored(path)
-
-        def is_dvcfile_and_not_ignored(root, file):
-            return is_valid_filename(file) and not is_ignored(
-                f"{root}{sep}{file}"
-            )
-
-        def is_out_or_ignored(root, directory):
-            dir_path = f"{root}{sep}{directory}"
-            # trailing slash needed to check if a directory is gitignored
-            return dir_path in outs or is_ignored(f"{dir_path}{sep}")
-
-        walk_iter = self.repo.dvcignore.walk(self.fs, self.repo.root_dir)
-        if logger.isEnabledFor(logging.TRACE):  # type: ignore[attr-defined]
-            walk_iter = log_walk(walk_iter)
-
-        for root, dirs, files in walk_iter:
-            dvcfile_filter = partial(is_dvcfile_and_not_ignored, root)
-            for file in filter(dvcfile_filter, files):
-                file_path = self.fs.path.join(root, file)
-                try:
-                    new_stages = self.load_file(file_path)
-                except DvcException as exc:
-                    if onerror:
-                        onerror(relpath(file_path), exc)
-                        continue
-                    raise
-
-                yield from new_stages
-                outs.update(
-                    out.fspath
-                    for stage in new_stages
-                    for out in stage.outs
-                    if out.protocol == "local"
-                )
-            dirs[:] = [d for d in dirs if not is_out_or_ignored(root, d)]
-
-    def collect_repo(self, onerror: Callable[[str, Exception], None] = None):
-        return list(self._collect_repo(onerror))
-
-    def _load_file(self, path):
-        from dvc.dvcfile import Dvcfile
-        from dvc.stage.loader import SingleStageLoader, StageLoader
-
-        path = self._get_filepath(path)
-        dvcfile = Dvcfile(self.repo, path)
-        stages = dvcfile.stages
-
-        if isinstance(stages, SingleStageLoader):
-            stages_ = [stages[None]]
-        else:
-            assert isinstance(stages, StageLoader)
-            keys = self._get_keys(stages)
-            stages_ = [stages[key] for key in keys]
-        return stages_, dvcfile
-
-    def _collect_all_from_repo(  # noqa: C901
-        self, onerror: Callable[[str, Exception], None] = None
-    ):
-        """Collects all of the stages present in the DVC repo.
-
-        Args:
-            onerror (optional): callable that will be called with two args:
-                the filepath whose collection failed and the exc instance.
-                It can report the error to continue with the collection
-                (and, skip failed ones), or raise the exception to abort
-                the collection.
-        """
-        from dvc.dvcfile import is_valid_filename
-        from dvc.fs import LocalFileSystem
-
-        scm = self.repo.scm
-        sep = self.fs.sep
-        outs: Set[str] = set()
-
-        is_local_fs = isinstance(self.fs, LocalFileSystem)
-
-        def is_ignored(path):
-            # apply only for the local fs
-            return is_local_fs and scm.is_ignored(path)
-
-        def is_dvcfile_and_not_ignored(root, file):
-            return is_valid_filename(file) and not is_ignored(
-                f"{root}{sep}{file}"
-            )
-
-        def is_out_or_ignored(root, directory):
-            dir_path = f"{root}{sep}{directory}"
-            # trailing slash needed to check if a directory is gitignored
-            return dir_path in outs or is_ignored(f"{dir_path}{sep}")
-
-        walk_iter = self.repo.dvcignore.walk(self.fs, self.repo.root_dir)
-        if logger.isEnabledFor(logging.TRACE):  # type: ignore[attr-defined]
-            walk_iter = log_walk(walk_iter)
-
-        stages = []
-        metrics = {}
-        plots = {}
-        params = {}
-
-        for root, dirs, files in walk_iter:
-            dvcfile_filter = partial(is_dvcfile_and_not_ignored, root)
-            for file in filter(dvcfile_filter, files):
-                file_path = self.fs.path.join(root, file)
-                try:
-                    new_stages, dvcfile = self._load_file(file_path)
-                except DvcException as exc:
-                    if onerror:
-                        onerror(relpath(file_path), exc)
-                        continue
-                    raise
-
-                stages.extend(new_stages)
-                if new_metrics := dvcfile.metrics:
-                    metrics[file_path] = new_metrics
-                if new_plots := dvcfile.plots:
-                    plots[file_path] = new_plots
-                if new_params := dvcfile.params:
-                    params[file_path] = new_params
-
-                outs.update(
-                    out.fspath
-                    for stage in new_stages
-                    for out in stage.outs
-                    if out.protocol == "local"
-                )
-            dirs[:] = [d for d in dirs if not is_out_or_ignored(root, d)]
-        return stages, metrics, plots, params

tests/func/test_remove.py

@@ -26,7 +26,7 @@ def test_remove(tmp_dir, scm, dvc, run_copy, remove_outs):
     for stage in [stage1, stage2, stage3]:
         dvc.remove(stage.addressing, outs=remove_outs)
         out_exists = (out.exists for out in stage.outs)
-        assert stage not in dvc.stage.collect_repo()
+        assert stage not in dvc.index.stages
         if remove_outs:
             assert not any(out_exists)
         else:

tests/func/test_repo_index.py

@@ -12,11 +12,9 @@ def test_index(tmp_dir, scm, dvc, run_copy):
     stage2 = run_copy("foo", "bar", name="copy-foo-bar")
     tmp_dir.commit([s.outs[0].fspath for s in (stage1, stage2)], msg="add")
 
-    index = Index(dvc)
+    index = Index.from_repo(dvc)
-    assert index.fs == dvc.fs
 
     assert set(index.stages) == {stage1, stage2}
-
     assert index.outs_graph
     assert index.graph
     assert isinstance(index.outs_trie, Trie)
@@ -70,7 +68,7 @@ def test_deps_outs_getters(tmp_dir, dvc, run_copy_metrics):
         name="copy-metrics2",
     )
 
-    index = Index(dvc)
+    index = Index.from_repo(dvc)
 
     stages = [foo_stage, run_stage1, run_stage2]
     (metrics,) = run_stage1.outs
@@ -88,16 +86,16 @@ def test_deps_outs_getters(tmp_dir, dvc, run_copy_metrics):
     assert outputs_equal(index.params, [params])
 
 
-def test_add_update(dvc):
+def test_update(dvc):
-    """Test that add/update overwrites existing stages with the new ones.
+    """Test that update overwrites existing stages with the new ones.
 
     The old stages and the new ones might have same hash, so we are
     making sure that the old stages were removed and replaced by new ones
     using `id`/`is` checks.
     """
-    index = Index(dvc)
+    index = Index.from_repo(dvc)
     new_stage = Stage(dvc, path="path1")
-    new_index = index.add(new_stage)
+    new_index = index.update({new_stage})
 
     assert not index.stages
     assert new_index.stages == [new_stage]
@@ -151,14 +149,14 @@ def get_index(dvc, rev):
         if rev != "workspace":
             assert next(brancher) == "workspace"
         next(brancher)
-    return Index(dvc)
+    return Index.from_repo(dvc)
 
 
 @pytest.mark.parametrize("rev", ["workspace", "HEAD"])
 def test_used_objs(tmp_dir, scm, dvc, run_copy, rev):
     from dvc_data.hashfile.hash_info import HashInfo
 
-    dvc.config["core"]["autostage"] = True
+    dvc.scm_context.autostage = True
     tmp_dir.dvc_gen({"dir": {"subdir": {"file": "file"}}, "foo": "foo"})
     run_copy("foo", "bar", name="copy-foo-bar")
     scm.commit("commit")
@@ -196,7 +194,7 @@ def test_view_granular_dir(tmp_dir, scm, dvc, run_copy):
         {"dir": {"subdir": {"in-subdir": "in-subdir"}, "in-dir": "in-dir"}},
         commit="init",
     )
-    index = Index(dvc)
+    index = Index.from_repo(dvc)
 
     # view should include the specific target, parent dirs, and children
     # view should exclude any siblings of the target
@@ -218,7 +216,7 @@ def test_view_stage_filter(tmp_dir, scm, dvc, run_copy):
     (stage1,) = tmp_dir.dvc_gen("foo", "foo")
     stage2 = run_copy("foo", "bar", name="copy-foo-bar")
     tmp_dir.commit([s.outs[0].fspath for s in (stage1, stage2)], msg="add")
-    index = Index(dvc)
+    index = Index.from_repo(dvc)
 
     view = index.targets_view(None)
     assert set(view.stages) == {stage1, stage2}
@@ -239,7 +237,7 @@ def test_view_outs_filter(tmp_dir, scm, dvc, run_copy):
     (stage1,) = tmp_dir.dvc_gen("foo", "foo")
     stage2 = run_copy("foo", "bar", name="copy-foo-bar")
     tmp_dir.commit([s.outs[0].fspath for s in (stage1, stage2)], msg="add")
-    index = Index(dvc)
+    index = Index.from_repo(dvc)
 
     view = index.targets_view(None, outs_filter=lambda o: o.def_path == "foo")
     assert set(view.stages) == {stage1, stage2}
@@ -252,7 +250,7 @@ def test_view_combined_filter(tmp_dir, scm, dvc, run_copy):
     (stage1,) = tmp_dir.dvc_gen("foo", "foo")
     stage2 = run_copy("foo", "bar", name="copy-foo-bar")
     tmp_dir.commit([s.outs[0].fspath for s in (stage1, stage2)], msg="add")
-    index = Index(dvc)
+    index = Index.from_repo(dvc)
 
     view = index.targets_view(
         None,
@@ -271,3 +269,24 @@ def test_view_combined_filter(tmp_dir, scm, dvc, run_copy):
     assert {out.fs_path for out in view.outs} == {
         out.fs_path for out in stage2.outs
     }
+
+
+def test_with_gitignore(tmp_dir, dvc, scm):
+    (stage,) = tmp_dir.dvc_gen({"data": {"foo": "foo", "bar": "bar"}})
+
+    index = Index.from_repo(dvc)
+    assert index.stages == [stage]
+
+    scm.ignore(stage.path)
+    scm._reset()
+
+    index = Index.from_repo(dvc)
+    assert not index.stages
+
+
+def test_ignored_dir_unignored_pattern(tmp_dir, dvc, scm):
+    tmp_dir.gen({".gitignore": "data/**\n!data/**/\n!data/**/*.dvc"})
+    scm.add([".gitignore"])
+    (stage,) = tmp_dir.dvc_gen({"data/raw/tracked.csv": "5,6,7,8"})
+    index = Index.from_repo(dvc)
+    assert index.stages == [stage]

tests/func/test_stage.py

@@ -239,13 +239,6 @@ def test_parent_repo_collect_stages(tmp_dir, scm, dvc):
     assert deep_subrepo_stages != []
 
 
-def test_collect_repo_ignored_dir_unignored_pattern(tmp_dir, dvc, scm):
-    tmp_dir.gen({".gitignore": "data/**\n!data/**/\n!data/**/*.dvc"})
-    scm.add([".gitignore"])
-    (stage,) = tmp_dir.dvc_gen({"data/raw/tracked.csv": "5,6,7,8"})
-    assert dvc.stage.collect_repo() == [stage]
-
-
 @pytest.mark.parametrize("with_deps", (False, True))
 def test_collect_symlink(tmp_dir, dvc, with_deps):
     tmp_dir.gen({"data": {"foo": "foo contents"}})

tests/func/test_stage_load.py

@@ -391,7 +391,7 @@ def test_collect_optimization(tmp_dir, dvc, mocker):
     # Forget cached stages and graph and error out on collection
     dvc._reset()
     mocker.patch(
-        "dvc.repo.index.Index.stages",
+        "dvc.repo.Repo.index",
         property(raiser(Exception("Should not collect"))),
     )
 
@@ -406,7 +406,7 @@ def test_collect_optimization_on_stage_name(tmp_dir, dvc, mocker, run_copy):
     # Forget cached stages and graph and error out on collection
     dvc._reset()
     mocker.patch(
-        "dvc.repo.index.Index.stages",
+        "dvc.repo.Repo.index",
         property(raiser(Exception("Should not collect"))),
     )
 
@@ -431,17 +431,6 @@ def test_collect_repo_callback(tmp_dir, dvc, mocker):
     assert isinstance(exc, YAMLValidationError)
 
 
-def test_gitignored_collect_repo(tmp_dir, dvc, scm):
-    (stage,) = tmp_dir.dvc_gen({"data": {"foo": "foo", "bar": "bar"}})
-
-    assert dvc.stage.collect_repo() == [stage]
-
-    scm.ignore(stage.path)
-    scm._reset()
-
-    assert not dvc.stage.collect_repo()
-
-
 def test_gitignored_file_try_collect_granular_for_data_files(
     tmp_dir, dvc, scm
 ):

tests/unit/command/test_status.py

@@ -76,12 +76,14 @@ def test_status_quiet(dvc, mocker, caplog, capsys, status, ret):
 
 
 def test_status_empty(dvc, mocker, capsys):
+    from dvc.repo.index import Index
+
     cli_args = parse_args(["status"])
     assert cli_args.func == CmdDataStatus
 
     cmd = cli_args.func(cli_args)
 
-    spy = mocker.spy(cmd.repo.stage, "_collect_all_from_repo")
+    spy = mocker.spy(Index, "from_repo")
 
     assert cmd.run() == 0
 
@@ -100,6 +102,8 @@ def test_status_empty(dvc, mocker, capsys):
     ],
 )
 def test_status_up_to_date(dvc, mocker, capsys, cloud_opts, expected_message):
+    from dvc.repo.index import Index
+
     cli_args = parse_args(["status", *cloud_opts])
     assert cli_args.func == CmdDataStatus
 
@@ -107,12 +111,8 @@ def test_status_up_to_date(dvc, mocker, capsys, cloud_opts, expected_message):
 
     mocker.patch.dict(cmd.repo.config, {"core": {"remote": "default"}})
     mocker.patch.object(cmd.repo, "status", autospec=True, return_value={})
-    mocker.patch.object(
+    mocker.patch("dvc.repo.Repo.index", return_value=Index(dvc, [object()]))
-        cmd.repo.stage,
+    cmd.repo._reset = mocker.Mock()
-        "_collect_all_from_repo",
-        return_value=[[object()], {}, {}, {}],
-        autospec=True,
-    )
 
     assert cmd.run() == 0
     captured = capsys.readouterr()