repro: Move pull logic to inside Stage.

iterative · May 16, 2023 · 8ecd5ab · 8ecd5ab
1 parent e69deaa
commit 8ecd5ab
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 10 deletions.
diff --git a/dvc/repo/reproduce.py b/dvc/repo/reproduce.py
@@ -194,10 +194,6 @@ def _reproduce_stages(  # noqa: C901
             )
 
         try:
-            if kwargs.get("pull") and stage.changed():
-                logger.debug("Pulling %s", stage.addressing)
-                stage.repo.pull(stage.addressing, allow_missing=True)
-
             ret = _reproduce_stage(stage, **kwargs)
 
             if len(ret) == 0:

diff --git a/dvc/stage/__init__.py b/dvc/stage/__init__.py
@@ -584,7 +584,7 @@ def add_outs(  # noqa: C901
             raise CacheLinkError(link_failures)
 
     @rwlocked(read=["deps", "outs"])
-    def run(
+    def run(  # noqa: C901
         self,
         dry=False,
         no_commit=False,
@@ -596,15 +596,20 @@ def run(
         if (self.cmd or self.is_import) and not self.frozen and not dry:
             self.remove_outs(ignore_remove=False, force=False)
 
-        if (not self.frozen and self.is_import) or self.is_partial_import:
+        if (
+            self.is_import and (not self.frozen or kwargs.get("pull"))
+        ) or self.is_partial_import:
             self._sync_import(dry, force, kwargs.get("jobs", None), no_download)
         elif not self.frozen and self.cmd:
             self._run_stage(dry, force, **kwargs)
-        else:
+        elif kwargs.get("pull"):
+            logger.info("Pulling data for %s", self)
+            self.repo.pull(self.addressing, jobs=kwargs.get("jobs", None))
+            self.checkout()
+        elif not dry:
             args = ("outputs", "frozen ") if self.frozen else ("data sources", "")
             logger.info("Verifying %s in %s%s", *args, self)
-            if not dry:
-                self._check_missing_outputs()
+            self._check_missing_outputs()
 
         if not dry:
             if kwargs.get("checkpoint_func", None) or no_download:

diff --git a/tests/func/test_run_cache.py b/tests/func/test_run_cache.py
@@ -182,7 +182,7 @@ def test_restore_pull(tmp_dir, dvc, run_copy, mocker, local_remote):
 
     mock_restore.assert_called_once_with(stage, pull=True, dry=False)
     mock_run.assert_not_called()
-    assert mock_checkout.call_count == 3
+    assert mock_checkout.call_count == 2
     assert (tmp_dir / "bar").exists()
     assert not (tmp_dir / "foo").unlink()
     assert (tmp_dir / LOCK_FILE).exists()