diff --git a/dvc/dvcfile.py b/dvc/dvcfile.py index 1a150e8c9..c84ff8574 100644 --- a/dvc/dvcfile.py +++ b/dvc/dvcfile.py @@ -162,6 +162,9 @@ class FileMixin: def dump(self, stage, **kwargs): raise NotImplementedError + def dump_stages(self, stages, **kwargs): + raise NotImplementedError + def merge(self, ancestor, other, allowed=None): raise NotImplementedError @@ -198,6 +201,13 @@ class SingleStageFile(FileMixin): dump_yaml(self.path, serialize.to_single_stage_file(stage, **kwargs)) self.repo.scm_context.track_file(self.relpath) + def dump_stages(self, stages, **kwargs) -> None: + if not stages: + return + + assert len(stages) == 1, "SingleStageFile can only dump one stage." + return self.dump(stages[0], **kwargs) + def remove_stage(self, stage): # noqa: ARG002 self.remove() @@ -228,17 +238,27 @@ class ProjectFile(FileMixin): def dump(self, stage, update_pipeline=True, update_lock=True, **kwargs): """Dumps given stage appropriately in the dvcfile.""" + return self.dump_stages( + [stage], update_pipeline=update_pipeline, update_lock=update_lock, **kwargs + ) + + def dump_stages(self, stages, update_pipeline=True, update_lock=True, **kwargs): from dvc.stage import PipelineStage - assert isinstance(stage, PipelineStage) + if not stages: + return + + for stage in stages: + assert isinstance(stage, PipelineStage) + if self.verify: check_dvcfile_path(self.repo, self.path) - if update_pipeline and not stage.is_data_source: - self._dump_pipeline_file(stage) + if update_pipeline: + self._dump_pipeline_file(stages) if update_lock: - self._dump_lockfile(stage, **kwargs) + self._dump_lockfile(stages, **kwargs) def dump_dataset(self, dataset): with modify_yaml(self.path, fs=self.repo.fs) as data: @@ -260,32 +280,37 @@ class ProjectFile(FileMixin): raw.append(dataset) self.repo.scm_context.track_file(self.relpath) - def _dump_lockfile(self, stage, **kwargs): - self._lockfile.dump(stage, **kwargs) + def _dump_lockfile(self, stages, **kwargs): + self._lockfile.dump_stages(stages, **kwargs) @staticmethod def _check_if_parametrized(stage, action: str = "dump") -> None: if stage.raw_data.parametrized: raise ParametrizedDumpError(f"cannot {action} a parametrized {stage}") - def _dump_pipeline_file(self, stage): - self._check_if_parametrized(stage) - stage_data = serialize.to_pipeline_file(stage) + def _dump_pipeline_file(self, stages): + stages = stages if isinstance(stages, list) else [stages] + if not stages: + return + + for stage in stages: + self._check_if_parametrized(stage) with modify_yaml(self.path, fs=self.repo.fs) as data: if not data: logger.info("Creating '%s'", self.relpath) data["stages"] = data.get("stages", {}) - existing_entry = stage.name in data["stages"] - action = "Modifying" if existing_entry else "Adding" - logger.info("%s stage '%s' in '%s'", action, stage.name, self.relpath) - - if existing_entry: - orig_stage_data = data["stages"][stage.name] - apply_diff(stage_data[stage.name], orig_stage_data) - else: - data["stages"].update(stage_data) + for stage in stages: + stage_data = serialize.to_pipeline_file(stage) + existing_entry = stage.name in data["stages"] + action = "Modifying" if existing_entry else "Adding" + logger.info("%s stage '%s' in '%s'", action, stage.name, self.relpath) + if existing_entry: + orig_stage_data = data["stages"][stage.name] + apply_diff(stage_data[stage.name], orig_stage_data) + else: + data["stages"].update(stage_data) self.repo.scm_context.track_file(self.relpath) @@ -399,9 +424,12 @@ class Lockfile(FileMixin): data.setdefault("stages", {}) self.repo.scm_context.track_file(self.relpath) - def dump(self, stage, **kwargs): - stage_data = serialize.to_lockfile(stage, **kwargs) + def dump_stages(self, stages, **kwargs): + if not stages: + return + is_modified = False + log_updated = False with modify_yaml(self.path, fs=self.repo.fs) as data: if not data: data.update({"schema": "2.0"}) @@ -409,17 +437,24 @@ class Lockfile(FileMixin): logger.info("Generating lock file '%s'", self.relpath) data["stages"] = data.get("stages", {}) - modified = data["stages"].get(stage.name, {}) != stage_data.get( - stage.name, {} - ) - if modified: - logger.info("Updating lock file '%s'", self.relpath) + for stage in stages: + stage_data = serialize.to_lockfile(stage, **kwargs) + modified = data["stages"].get(stage.name, {}) != stage_data.get( + stage.name, {} + ) + if modified: + is_modified = True + if not log_updated: + logger.info("Updating lock file '%s'", self.relpath) + log_updated = True + data["stages"].update(stage_data) - data["stages"].update(stage_data) - - if modified: + if is_modified: self.repo.scm_context.track_file(self.relpath) + def dump(self, stage, **kwargs): + self.dump_stages([stage], **kwargs) + def remove_stage(self, stage): if not self.exists(): return diff --git a/dvc/output.py b/dvc/output.py index 48acf499b..6a615549a 100644 --- a/dvc/output.py +++ b/dvc/output.py @@ -1452,25 +1452,43 @@ class Output: self.remote = other.remote self.can_push = other.can_push - def merge_version_meta(self, other: "Output"): + def _get_versioned_meta( + self, + ) -> Optional[ + tuple["HashInfo", Optional["Meta"], Optional[Union["HashFile", "Tree"]]] + ]: + if self.files is not None or ( + self.meta is not None and self.meta.version_id is not None + ): + old_obj = self.obj if self.obj is not None else self.get_obj() + return self.hash_info, self.meta, old_obj + return None + + def merge_version_meta( + self, + old_hi: "HashInfo", + old_meta: Optional["Meta"], + old_obj: Optional[Union["HashFile", "Tree"]], + ): """Merge version meta for files which are unchanged from other.""" if not self.hash_info: return if self.hash_info.isdir: - return self._merge_dir_version_meta(other) - if self.hash_info != other.hash_info: + return self._merge_dir_version_meta(old_hi, old_obj) + if self.hash_info != old_hi: return - self.meta = other.meta + self.meta = old_meta - def _merge_dir_version_meta(self, other: "Output"): + def _merge_dir_version_meta( + self, old_hi: "HashInfo", old_obj: Optional[Union["HashFile", "Tree"]] + ): from dvc_data.hashfile.tree import update_meta - if not self.obj or not other.hash_info.isdir: + if not self.obj or not old_hi.isdir: return - other_obj = other.obj if other.obj is not None else other.get_obj() assert isinstance(self.obj, Tree) - assert isinstance(other_obj, Tree) - updated = update_meta(self.obj, other_obj) + assert isinstance(old_obj, Tree) + updated = update_meta(self.obj, old_obj) assert updated.hash_info == self.obj.hash_info self.obj = updated self.files = updated.as_list(with_meta=True) diff --git a/dvc/repo/commit.py b/dvc/repo/commit.py index cea727f20..986af04b1 100644 --- a/dvc/repo/commit.py +++ b/dvc/repo/commit.py @@ -1,3 +1,4 @@ +from itertools import groupby from typing import TYPE_CHECKING from dvc import prompt @@ -52,29 +53,38 @@ def commit( data_only=False, relink=True, ): - stages_info = [ - info - for info in self.stage.collect_granular( - target, with_deps=with_deps, recursive=recursive - ) - if not data_only or info.stage.is_data_source - ] - for stage_info in stages_info: - stage = stage_info.stage - if force: - stage.save(allow_missing=allow_missing) - else: - changes = stage.changed_entries() - if any(changes): - prompt_to_commit(stage, changes, force=force) + committed_stages = [] + groups = groupby( + [ + info + for info in self.stage.collect_granular( + target, with_deps=with_deps, recursive=recursive + ) + if not data_only or info.stage.is_data_source + ], + key=lambda info: info.stage.dvcfile, + ) + + for dvcfile, stages_info_group in groups: + to_dump = [] + for stage_info in stages_info_group: + stage = stage_info.stage + if force: stage.save(allow_missing=allow_missing) - stage.commit( - filter_info=stage_info.filter_info, - allow_missing=allow_missing, - relink=relink, - ) - stage.dump(update_pipeline=False) - return [s.stage for s in stages_info] + else: + changes = stage.changed_entries() + if any(changes): + prompt_to_commit(stage, changes, force=force) + stage.save(allow_missing=allow_missing) + stage.commit( + filter_info=stage_info.filter_info, + allow_missing=allow_missing, + relink=relink, + ) + to_dump.append(stage) + dvcfile.dump_stages(to_dump, update_pipeline=False) + committed_stages.extend(to_dump) + return committed_stages @locked diff --git a/dvc/stage/__init__.py b/dvc/stage/__init__.py index 92835fe4c..fa1270d6c 100644 --- a/dvc/stage/__init__.py +++ b/dvc/stage/__init__.py @@ -509,34 +509,20 @@ class Stage(params.StageParams): if not allow_missing: raise - def get_versioned_outs(self) -> dict[str, "Output"]: - from .exceptions import StageFileDoesNotExistError, StageNotFound - - try: - old = self.reload() - except (StageFileDoesNotExistError, StageNotFound): - return {} - - return { - out.def_path: out - for out in old.outs - if out.files is not None - or (out.meta is not None and out.meta.version_id is not None) - } - def save_outs(self, allow_missing: bool = False): from dvc.output import OutputDoesNotExistError - old_versioned_outs = self.get_versioned_outs() for out in self.outs: + # old state just before saving so to merge them later + old_state = out._get_versioned_meta() try: out.save() except OutputDoesNotExistError: if not allow_missing: raise - if old_out := old_versioned_outs.get(out.def_path): - out.merge_version_meta(old_out) + if old_state: + out.merge_version_meta(*old_state) def ignore_outs(self) -> None: for out in self.outs: @@ -579,8 +565,9 @@ class Stage(params.StageParams): from dvc.output import OutputDoesNotExistError link_failures = [] - old_versioned_outs = self.get_versioned_outs() for out in self.filter_outs(filter_info): + # old state just before saving so to merge them later + old_state = out._get_versioned_meta() try: out.add(filter_info, **kwargs) except (FileNotFoundError, OutputDoesNotExistError): @@ -589,8 +576,8 @@ class Stage(params.StageParams): except CacheLinkError: link_failures.append(filter_info or out.fs_path) - if old_out := old_versioned_outs.get(out.def_path): - out.merge_version_meta(old_out) + if old_state: + out.merge_version_meta(*old_state) if link_failures: raise CacheLinkError(link_failures) diff --git a/tests/func/test_commit.py b/tests/func/test_commit.py index 2c91dab61..645da15d7 100644 --- a/tests/func/test_commit.py +++ b/tests/func/test_commit.py @@ -4,7 +4,7 @@ import textwrap import pytest from dvc.dependency.base import DependencyDoesNotExistError -from dvc.dvcfile import PROJECT_FILE +from dvc.dvcfile import PROJECT_FILE, Lockfile, ProjectFile, SingleStageFile from dvc.fs import localfs from dvc.output import OutputDoesNotExistError from dvc.stage.exceptions import StageCommitError @@ -325,3 +325,51 @@ def test_commit_dos2unix(tmp_dir, dvc): dvc.commit("foo.dvc", force=True) content = (tmp_dir / "foo.dvc").read_text() assert "hash: md5" in content + + +def test_commit_multiple_files(tmp_dir, dvc, mocker): + tmp_dir.gen({"foo": "foo", "bar": "bar"}) + stages = dvc.add(["foo", "bar"], no_commit=True) + test1_stage = dvc.stage.add(name="test", cmd="echo test", deps=["foo"]) + test2_stage = dvc.stage.add(name="test2", cmd="echo test2", deps=["foo"]) + + subdir = tmp_dir / "subdir" + subdir.mkdir() + with subdir.chdir(): + bar_relpath = os.path.relpath(tmp_dir / "bar", subdir) + test3_stage = dvc.stage.add(name="test3", cmd="echo test3", deps=[bar_relpath]) + + pointerfile_spy = mocker.spy(SingleStageFile, "dump_stages") + projectfile_spy = mocker.spy(ProjectFile, "dump_stages") + lockfile_spy = mocker.spy(Lockfile, "dump_stages") + + assert set(dvc.commit(force=True)) == { + *stages, + test1_stage, + test2_stage, + test3_stage, + } + pointerfile_spy.assert_has_calls( + [ + mocker.call(stages[0].dvcfile, [stages[0]], update_pipeline=False), + mocker.call(stages[1].dvcfile, [stages[1]], update_pipeline=False), + ], + any_order=True, + ) + projectfile_spy.assert_has_calls( + [ + mocker.call( + test1_stage.dvcfile, [test1_stage, test2_stage], update_pipeline=False + ), + mocker.call(test3_stage.dvcfile, [test3_stage], update_pipeline=False), + ], + any_order=True, + ) + lockfile_spy.assert_has_calls( + [ + mocker.call(test1_stage.dvcfile._lockfile, [test1_stage, test2_stage]), + mocker.call(test3_stage.dvcfile._lockfile, [test3_stage]), + ], + any_order=True, + ) + assert dvc.status() == {} diff --git a/tests/unit/test_dvcfile.py b/tests/unit/test_dvcfile.py index 11641388c..11a000641 100644 --- a/tests/unit/test_dvcfile.py +++ b/tests/unit/test_dvcfile.py @@ -90,6 +90,39 @@ def test_dump_stage(tmp_dir, dvc): assert list(dvcfile.stages.values()) == [stage] +def test_dump_multiple_pipeline_stages(tmp_dir, dvc): + stage1 = PipelineStage(dvc, cmd="cmd1", name="stage1", path="dvc.yaml") + stage2 = PipelineStage(dvc, cmd="cmd2", name="stage2", path="dvc.yaml") + dvcfile = load_file(dvc, "dvc.yaml") + + dvcfile.dump_stages([stage1, stage2], update_lock=False, update_pipeline=False) + assert not (tmp_dir / LOCK_FILE).exists() + assert not (tmp_dir / PROJECT_FILE).exists() + + dvcfile.dump_stages([stage1, stage2], update_pipeline=False) + assert not (tmp_dir / PROJECT_FILE).exists() + assert (tmp_dir / LOCK_FILE).parse() == { + "schema": "2.0", + "stages": {"stage1": {"cmd": "cmd1"}, "stage2": {"cmd": "cmd2"}}, + } + + dvcfile.dump_stages([stage1, stage2], update_lock=False) + assert (tmp_dir / PROJECT_FILE).parse() == { + "stages": {"stage1": {"cmd": "cmd1"}, "stage2": {"cmd": "cmd2"}} + } + + +def test_dump_stages_single_stage(tmp_dir, dvc): + stage = dvc.stage.create( + fname="foo.dvc", outs=["out"], deps=["dep"], single_stage=True + ) + stage.dvcfile.dump_stages([stage]) + assert (tmp_dir / "foo.dvc").parse() == { + "deps": [{"hash": "md5", "path": "dep"}], + "outs": [{"hash": "md5", "path": "out"}], + } + + @pytest.mark.parametrize("file", ["stage.dvc", "dvc.yaml"]) def test_stage_load_file_exists_but_dvcignored(tmp_dir, dvc, scm, file): (tmp_dir / file).write_text("")