diff --git a/dvc/repo/__init__.py b/dvc/repo/__init__.py index bd8edfb11..831f53c41 100644 --- a/dvc/repo/__init__.py +++ b/dvc/repo/__init__.py @@ -3,7 +3,7 @@ import os from collections import defaultdict from contextlib import contextmanager from functools import wraps -from typing import TYPE_CHECKING, Callable, Optional +from typing import TYPE_CHECKING, Callable, Iterable, Optional from funcy import cached_property @@ -18,6 +18,7 @@ if TYPE_CHECKING: from dvc.fs import FileSystem from dvc.repo.scm_context import SCMContext from dvc.scm import Base + from dvc.stage import Stage logger = logging.getLogger(__name__) @@ -241,7 +242,9 @@ class Repo: return Index.from_repo(self) - def ensure_graph_correctness_with(self, stages=None, callback=None): + def check_graph( + self, stages: Iterable["Stage"] = None, callback: Callable = None + ) -> None: if not getattr(self, "_skip_graph_checks", False): new = self.index.update(stages) if callable(callback): diff --git a/dvc/repo/add.py b/dvc/repo/add.py index be53fac56..7f139bcd0 100644 --- a/dvc/repo/add.py +++ b/dvc/repo/add.py @@ -180,7 +180,7 @@ 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. - repo.ensure_graph_correctness_with( + repo.check_graph( stages=stages, callback=lambda: status.update("Checking graph") ) diff --git a/dvc/repo/imp_url.py b/dvc/repo/imp_url.py index 144a84d76..40773f4bc 100644 --- a/dvc/repo/imp_url.py +++ b/dvc/repo/imp_url.py @@ -1,17 +1,21 @@ import os +from typing import TYPE_CHECKING from dvc.exceptions import InvalidArgumentError, OutputDuplicationError from dvc.repo.scm_context import scm_context from dvc.utils import relpath, resolve_output, resolve_paths from dvc.utils.fs import path_isin +if TYPE_CHECKING: + from . import Repo + from . import locked @locked @scm_context def imp_url( # noqa: C901 - self, + self: "Repo", url, out=None, fname=None, @@ -78,7 +82,7 @@ def imp_url( # noqa: C901 dvcfile.remove() try: - self.ensure_graph_correctness_with(stages={stage}) + self.check_graph(stages={stage}) except OutputDuplicationError as exc: raise OutputDuplicationError( # noqa: B904 exc.output, set(exc.stages) - {stage} diff --git a/dvc/repo/stage.py b/dvc/repo/stage.py index 48e221e0b..7163f924d 100644 --- a/dvc/repo/stage.py +++ b/dvc/repo/stage.py @@ -185,7 +185,7 @@ class StageLoad: check_stage_exists(self.repo, stage, stage.path) - self.repo.ensure_graph_correctness_with(stages={stage}) + self.repo.check_graph(stages={stage}) restore_fields(stage) return stage