diff --git a/dvc/annotations.py b/dvc/annotations.py new file mode 100644 index 000000000..eb3b856a0 --- /dev/null +++ b/dvc/annotations.py @@ -0,0 +1,32 @@ +from dataclasses import asdict, dataclass, field, fields +from typing import ClassVar, Dict, List, Optional + +from funcy import compact + + +@dataclass +class Annotation: + PARAM_DESC: ClassVar[str] = "desc" + PARAM_TYPE: ClassVar[str] = "type" + PARAM_LABELS: ClassVar[str] = "labels" + + desc: Optional[str] = None + type: Optional[str] = None + labels: List[str] = field(default_factory=list) + + def update(self, **kwargs) -> "Annotation": + for attr, value in kwargs.items(): + if value and hasattr(self, attr): + setattr(self, attr, value) + return self + + def to_dict(self) -> Dict[str, str]: + return compact(asdict(self)) + + +ANNOTATION_FIELDS = [field.name for field in fields(Annotation)] +ANNOTATION_SCHEMA = { + Annotation.PARAM_DESC: str, + Annotation.PARAM_TYPE: str, + Annotation.PARAM_LABELS: [str], +} diff --git a/dvc/commands/add.py b/dvc/commands/add.py index 1cc508496..091519bd7 100644 --- a/dvc/commands/add.py +++ b/dvc/commands/add.py @@ -8,6 +8,32 @@ from dvc.cli.utils import append_doc_link logger = logging.getLogger(__name__) +def _add_annotating_args(parser: argparse.ArgumentParser) -> None: + parser.add_argument( + "--desc", + type=str, + metavar="", + help=( + "User description of the data (optional). " + "This doesn't affect any DVC operations." + ), + ) + parser.add_argument( + "--label", + dest="labels", + type=str, + action="append", + metavar="", + help="Comma separated list of labels for the data (optional).", + ) + parser.add_argument( + "--type", + type=str, + metavar="", + help="Type of the data (optional).", + ) + + class CmdAdd(CmdBase): def run(self): from dvc.exceptions import ( @@ -28,6 +54,8 @@ class CmdAdd(CmdBase): glob=self.args.glob, desc=self.args.desc, out=self.args.out, + type=self.args.type, + labels=self.args.labels, remote=self.args.remote, to_remote=self.args.to_remote, jobs=self.args.jobs, @@ -109,15 +137,8 @@ def add_parser(subparsers, parent_parser): ), metavar="", ) - parser.add_argument( - "--desc", - type=str, - metavar="", - help=( - "User description of the data (optional). " - "This doesn't affect any DVC operations." - ), - ) + + _add_annotating_args(parser) parser.add_argument( "targets", nargs="+", help="Input files/directories to add." ).complete = completion.FILE diff --git a/dvc/commands/imp.py b/dvc/commands/imp.py index 8b7df1785..bfa75445c 100644 --- a/dvc/commands/imp.py +++ b/dvc/commands/imp.py @@ -21,6 +21,8 @@ class CmdImport(CmdBase): no_exec=self.args.no_exec, no_download=self.args.no_download, desc=self.args.desc, + type=self.args.type, + labels=self.args.labels, jobs=self.args.jobs, ) except DvcException: @@ -34,6 +36,8 @@ class CmdImport(CmdBase): def add_parser(subparsers, parent_parser): + from .add import _add_annotating_args + IMPORT_HELP = ( "Download file or directory tracked by DVC or by Git " "into the workspace, and track it." @@ -84,15 +88,6 @@ def add_parser(subparsers, parent_parser): help="Create .dvc file including target data hash value(s)" " but do not actually download the file(s).", ) - import_parser.add_argument( - "--desc", - type=str, - metavar="", - help=( - "User description of the data (optional). " - "This doesn't affect any DVC operations." - ), - ) import_parser.add_argument( "-j", "--jobs", @@ -103,4 +98,6 @@ def add_parser(subparsers, parent_parser): ), metavar="", ) + + _add_annotating_args(import_parser) import_parser.set_defaults(func=CmdImport) diff --git a/dvc/commands/imp_url.py b/dvc/commands/imp_url.py index add80884c..baa6e681a 100644 --- a/dvc/commands/imp_url.py +++ b/dvc/commands/imp_url.py @@ -21,6 +21,8 @@ class CmdImportUrl(CmdBase): remote=self.args.remote, to_remote=self.args.to_remote, desc=self.args.desc, + type=self.args.type, + labels=self.args.labels, jobs=self.args.jobs, ) except DvcException: @@ -35,6 +37,8 @@ class CmdImportUrl(CmdBase): def add_parser(subparsers, parent_parser): + from .add import _add_annotating_args + IMPORT_HELP = ( "Download or copy file from URL and take it under DVC control." ) @@ -103,13 +107,5 @@ def add_parser(subparsers, parent_parser): ), metavar="", ) - import_parser.add_argument( - "--desc", - type=str, - metavar="", - help=( - "User description of the data (optional). " - "This doesn't affect any DVC operations." - ), - ) + _add_annotating_args(import_parser) import_parser.set_defaults(func=CmdImportUrl) diff --git a/dvc/output.py b/dvc/output.py index fb77d080e..dd2dd929d 100644 --- a/dvc/output.py +++ b/dvc/output.py @@ -28,6 +28,7 @@ from dvc_data.hashfile.meta import Meta from dvc_data.hashfile.transfer import transfer as otransfer from dvc_objects.errors import ObjectFormatError +from .annotations import ANNOTATION_FIELDS, ANNOTATION_SCHEMA, Annotation from .fs import LocalFileSystem, RemoteMissingDepsError, Schemes, get_cloud_fs from .utils import relpath from .utils.fs import path_isin @@ -80,9 +81,9 @@ def loadd_from(stage, d_list): plot = d.pop(Output.PARAM_PLOT, False) persist = d.pop(Output.PARAM_PERSIST, False) checkpoint = d.pop(Output.PARAM_CHECKPOINT, False) - desc = d.pop(Output.PARAM_DESC, False) live = d.pop(Output.PARAM_LIVE, False) remote = d.pop(Output.PARAM_REMOTE, None) + annot = {field: d.pop(field, None) for field in ANNOTATION_FIELDS} ret.append( _get( stage, @@ -93,9 +94,9 @@ def loadd_from(stage, d_list): plot=plot, persist=persist, checkpoint=checkpoint, - desc=desc, live=live, remote=remote, + **annot, ) ) return ret @@ -187,7 +188,7 @@ def load_from_pipeline(stage, data, typ="outs"): Output.PARAM_PERSIST, Output.PARAM_CHECKPOINT, Output.PARAM_REMOTE, - Output.PARAM_DESC, + Annotation.PARAM_DESC, ], ) @@ -253,7 +254,6 @@ class Output: PARAM_PLOT_TITLE = "title" PARAM_PLOT_HEADER = "header" PARAM_PERSIST = "persist" - PARAM_DESC = "desc" PARAM_LIVE = "live" PARAM_LIVE_SUMMARY = "summary" PARAM_LIVE_HTML = "html" @@ -285,9 +285,12 @@ class Output: checkpoint=False, live=False, desc=None, + type=None, # pylint: disable=redefined-builtin + labels=None, remote=None, repo=None, ): + self.annot = Annotation(desc=desc, type=type, labels=labels or []) self.repo = stage.repo if not repo and stage else repo meta = Meta.from_dict(info or {}) # NOTE: when version_aware is not passed into get_cloud_fs, it will be @@ -336,7 +339,6 @@ class Output: self.persist = persist self.checkpoint = checkpoint self.live = live - self.desc = desc self.fs_path = self._parse_path(self.fs, fs_path) self.obj = None @@ -703,9 +705,7 @@ class Output: if self.IS_DEPENDENCY: return ret - if self.desc: - ret[self.PARAM_DESC] = self.desc - + ret.update(self.annot.to_dict()) if not self.use_cache: ret[self.PARAM_CACHE] = self.use_cache @@ -1121,22 +1121,26 @@ class Output: return bool(self.plot) or bool(self.live) -ARTIFACT_SCHEMA = { - **CHECKSUMS_SCHEMA, - Required(Output.PARAM_PATH): str, - Output.PARAM_PLOT: bool, - Output.PARAM_PERSIST: bool, - Output.PARAM_CHECKPOINT: bool, +META_SCHEMA = { Meta.PARAM_SIZE: int, Meta.PARAM_NFILES: int, Meta.PARAM_ISEXEC: bool, Meta.PARAM_VERSION_ID: str, } +ARTIFACT_SCHEMA = { + **CHECKSUMS_SCHEMA, + **META_SCHEMA, + Required(Output.PARAM_PATH): str, + Output.PARAM_PLOT: bool, + Output.PARAM_PERSIST: bool, + Output.PARAM_CHECKPOINT: bool, +} + SCHEMA = { **ARTIFACT_SCHEMA, + **ANNOTATION_SCHEMA, Output.PARAM_CACHE: bool, Output.PARAM_METRIC: Output.METRIC_SCHEMA, - Output.PARAM_DESC: str, Output.PARAM_REMOTE: str, } diff --git a/dvc/repo/add.py b/dvc/repo/add.py index 1bfb2c499..3c28316d0 100644 --- a/dvc/repo/add.py +++ b/dvc/repo/add.py @@ -264,6 +264,7 @@ def create_stages( outs=[out], external=external, ) - if kwargs.get("desc"): - stage.outs[0].desc = kwargs["desc"] + + out_obj = stage.outs[0] + out_obj.annot.update(**kwargs) yield stage diff --git a/dvc/repo/imp_url.py b/dvc/repo/imp_url.py index fc2f709a0..ae1d1a0f1 100644 --- a/dvc/repo/imp_url.py +++ b/dvc/repo/imp_url.py @@ -22,6 +22,8 @@ def imp_url( remote=None, to_remote=False, desc=None, + type=None, # pylint: disable=redefined-builtin + labels=None, jobs=None, ): from dvc.dvcfile import Dvcfile @@ -61,9 +63,8 @@ def imp_url( ) restore_fields(stage) - if desc: - stage.outs[0].desc = desc - + out_obj = stage.outs[0] + out_obj.annot.update(desc=desc, type=type, labels=labels) dvcfile = Dvcfile(self, stage.path) dvcfile.remove() diff --git a/dvc/schema.py b/dvc/schema.py index 67294bda8..6c44b2f9b 100644 --- a/dvc/schema.py +++ b/dvc/schema.py @@ -3,6 +3,7 @@ from collections.abc import Mapping from voluptuous import Any, Optional, Required, Schema from dvc import dependency, output +from dvc.annotations import Annotation from dvc.output import CHECKSUMS_SCHEMA, Output from dvc.parsing import DO_KWD, FOREACH_KWD, VARS_KWD from dvc.parsing.versions import SCHEMA_KWD, lockfile_version_schema @@ -50,7 +51,7 @@ OUT_PSTAGE_DETAILED_SCHEMA = { Output.PARAM_CACHE: bool, Output.PARAM_PERSIST: bool, Output.PARAM_CHECKPOINT: bool, - Output.PARAM_DESC: str, + Annotation.PARAM_DESC: str, Output.PARAM_REMOTE: str, } } diff --git a/dvc/stage/__init__.py b/dvc/stage/__init__.py index 85f686865..a0fa670b8 100644 --- a/dvc/stage/__init__.py +++ b/dvc/stage/__init__.py @@ -110,16 +110,13 @@ def restore_fields(stage): # will be used to restore comments later # noqa, pylint: disable=protected-access stage._stage_text = old._stage_text - stage.meta = old.meta stage.desc = old.desc - old_fields = {out.def_path: (out.desc, out.remote) for out in old.outs} - + old_fields = {out.def_path: (out.annot, out.remote) for out in old.outs} for out in stage.outs: - out_fields = old_fields.get(out.def_path, None) - if out_fields: - out.desc, out.remote = out_fields + if out_fields := old_fields.get(out.def_path, None): + out.annot, out.remote = out_fields class Stage(params.StageParams): diff --git a/dvc/stage/serialize.py b/dvc/stage/serialize.py index 2a1e15b05..866008354 100644 --- a/dvc/stage/serialize.py +++ b/dvc/stage/serialize.py @@ -6,7 +6,7 @@ from typing import TYPE_CHECKING, List, no_type_check from funcy import post_processing from dvc.dependency import ParamsDependency -from dvc.output import Output +from dvc.output import Annotation, Output from dvc.utils.collections import apply_diff from dvc.utils.serialize import parse_yaml_for_update @@ -27,7 +27,7 @@ PARAM_METRIC = Output.PARAM_METRIC PARAM_PLOT = Output.PARAM_PLOT PARAM_PERSIST = Output.PARAM_PERSIST PARAM_CHECKPOINT = Output.PARAM_CHECKPOINT -PARAM_DESC = Output.PARAM_DESC +PARAM_DESC = Annotation.PARAM_DESC PARAM_REMOTE = Output.PARAM_REMOTE DEFAULT_PARAMS_FILE = ParamsDependency.DEFAULT_PARAMS_FILE @@ -38,8 +38,8 @@ sort_by_path = partial(sorted, key=attrgetter("def_path")) @post_processing(OrderedDict) def _get_flags(out): - if out.desc: - yield PARAM_DESC, out.desc + if out.annot.desc: + yield PARAM_DESC, out.annot.desc if not out.use_cache: yield PARAM_CACHE, False if out.checkpoint: diff --git a/dvc/stage/utils.py b/dvc/stage/utils.py index abd630203..2e2b66d24 100644 --- a/dvc/stage/utils.py +++ b/dvc/stage/utils.py @@ -4,6 +4,7 @@ from typing import TYPE_CHECKING, Union from funcy import concat, first, lsplit, rpartial, without +from dvc.annotations import ANNOTATION_FIELDS from dvc.exceptions import InvalidArgumentError from dvc_data.hashfile.meta import Meta @@ -194,9 +195,9 @@ def compute_md5(stage): return dict_md5( d, exclude=[ + *ANNOTATION_FIELDS, stage.PARAM_LOCKED, # backward compatibility stage.PARAM_FROZEN, - Output.PARAM_DESC, Output.PARAM_METRIC, Output.PARAM_PERSIST, Output.PARAM_CHECKPOINT, diff --git a/tests/func/test_add.py b/tests/func/test_add.py index a072ab90d..863ba206e 100644 --- a/tests/func/test_add.py +++ b/tests/func/test_add.py @@ -12,6 +12,7 @@ import pytest import dvc as dvc_module import dvc_data +from dvc.annotations import Annotation from dvc.cli import main from dvc.dvcfile import DVC_FILE_SUFFIX from dvc.exceptions import ( @@ -903,6 +904,10 @@ def test_add_preserve_fields(tmp_dir, dvc): outs: - path: foo # out comment desc: out desc + type: mytype + labels: + - label1 + - label2 remote: testremote meta: some metadata """ @@ -916,6 +921,10 @@ def test_add_preserve_fields(tmp_dir, dvc): outs: - path: foo # out comment desc: out desc + type: mytype + labels: + - label1 + - label2 remote: testremote md5: acbd18db4cc2f85cedef654fccc4a4d8 size: 3 @@ -1161,3 +1170,18 @@ def test_add_ignore_duplicated_targets(tmp_dir, dvc, capsys): _, err = capsys.readouterr() assert len(stages) == 3 assert "ignoring duplicated targets: foo, bar" in err + + +def test_add_with_annotations(M, tmp_dir, dvc): + tmp_dir.gen("foo", "foo") + + annot = {"desc": "foo desc", "labels": ["l1", "l2"], "type": "t1"} + (stage,) = dvc.add("foo", **annot) + assert stage.outs[0].annot == Annotation(**annot) + assert (tmp_dir / "foo.dvc").parse() == M.dict(outs=[M.dict(**annot)]) + + # try to selectively update/overwrite some annotations + annot = {**annot, "type": "t2"} + (stage,) = dvc.add("foo", type="t2") + assert stage.outs[0].annot == Annotation(**annot) + assert (tmp_dir / "foo.dvc").parse() == M.dict(outs=[M.dict(**annot)]) diff --git a/tests/func/test_commit.py b/tests/func/test_commit.py index b1b40ce4e..b6575be6d 100644 --- a/tests/func/test_commit.py +++ b/tests/func/test_commit.py @@ -47,6 +47,10 @@ def test_commit_preserve_fields(tmp_dir, dvc): outs: - path: foo # out comment desc: out desc + type: mytype + labels: + - label1 + - label2 remote: testremote meta: some metadata """ @@ -61,6 +65,10 @@ def test_commit_preserve_fields(tmp_dir, dvc): outs: - path: foo # out comment desc: out desc + type: mytype + labels: + - label1 + - label2 remote: testremote md5: acbd18db4cc2f85cedef654fccc4a4d8 size: 3 diff --git a/tests/func/test_dvcfile.py b/tests/func/test_dvcfile.py index fdb4860f4..eb6880909 100644 --- a/tests/func/test_dvcfile.py +++ b/tests/func/test_dvcfile.py @@ -335,7 +335,7 @@ def test_dvcfile_dump_preserves_desc(tmp_dir, dvc, run_copy): (tmp_dir / dvcfile.path).dump(data) assert stage.desc == stage_desc - stage.outs[0].desc = out_desc + stage.outs[0].annot.desc = out_desc dvcfile.dump(stage) loaded = dvcfile._load()[0] assert loaded == data @@ -401,7 +401,7 @@ def test_dvcfile_load_dump_stage_with_desc_meta(tmp_dir, dvc): stage = dvc.stage.load_one(name="stage1") assert stage.meta == {"key1": "value1", "key2": "value2"} assert stage.desc == "stage desc" - assert stage.outs[0].desc == "bar desc" + assert stage.outs[0].annot.desc == "bar desc" # sanity check stage.dump() diff --git a/tests/func/test_import.py b/tests/func/test_import.py index cb0ba08b4..5f01d24d1 100644 --- a/tests/func/test_import.py +++ b/tests/func/test_import.py @@ -6,6 +6,7 @@ import pytest from funcy import first from scmrepo.git import Git +from dvc.annotations import Annotation from dvc.config import NoRemoteError from dvc.dvcfile import Dvcfile from dvc.exceptions import DownloadError, PathMissingError @@ -627,3 +628,21 @@ def test_parameterized_repo(tmp_dir, dvc, scm, erepo_dir, paths): "url": os.fspath(erepo_dir), "rev_lock": erepo_dir.scm.get_rev(), } + + +def test_import_with_annotations(M, tmp_dir, scm, dvc, erepo_dir): + with erepo_dir.chdir(): + erepo_dir.dvc_gen("foo", "foo content", commit="create foo") + + annot = {"desc": "foo desc", "labels": ["l1", "l2"], "type": "t1"} + stage = dvc.imp(os.fspath(erepo_dir), "foo", "foo", no_exec=True, **annot) + assert stage.outs[0].annot == Annotation(**annot) + assert (tmp_dir / "foo.dvc").parse() == M.dict(outs=[M.dict(**annot)]) + + # try to selectively update/overwrite some annotations + annot = {**annot, "type": "t2"} + stage = dvc.imp( + os.fspath(erepo_dir), "foo", "foo", no_exec=True, type="t2" + ) + assert stage.outs[0].annot == Annotation(**annot) + assert (tmp_dir / "foo.dvc").parse() == M.dict(outs=[M.dict(**annot)]) diff --git a/tests/func/test_import_url.py b/tests/func/test_import_url.py index 50dcf881c..55e13a2e9 100644 --- a/tests/func/test_import_url.py +++ b/tests/func/test_import_url.py @@ -5,6 +5,7 @@ from uuid import uuid4 import pytest +from dvc.annotations import Annotation from dvc.cli import main from dvc.dependency.base import DependencyDoesNotExistError from dvc.exceptions import InvalidArgumentError @@ -132,7 +133,7 @@ class TestImport(_TestImport): return False -def test_import_url_preserve_meta(tmp_dir, dvc): +def test_import_url_preserve_fields(tmp_dir, dvc): text = textwrap.dedent( """\ # top comment @@ -142,6 +143,10 @@ def test_import_url_preserve_meta(tmp_dir, dvc): outs: - path: bar # out comment desc: out desc + type: mytype + labels: + - label1 + - label2 meta: some metadata """ ) @@ -160,6 +165,10 @@ def test_import_url_preserve_meta(tmp_dir, dvc): outs: - path: bar # out comment desc: out desc + type: mytype + labels: + - label1 + - label2 md5: acbd18db4cc2f85cedef654fccc4a4d8 size: 3 meta: some metadata @@ -288,3 +297,27 @@ def test_import_url_no_download(tmp_dir, dvc, local_workspace): assert status["file.dvc"] == [ {"changed outs": {"file": "deleted"}}, ] + + +def test_imp_url_with_annotations(M, tmp_dir, dvc, local_workspace): + local_workspace.gen("foo", "foo") + annot = {"desc": "foo desc", "labels": ["l1", "l2"], "type": "t1"} + stage = dvc.imp_url( + "remote://workspace/foo", + os.fspath(tmp_dir / "foo"), + no_exec=True, + **annot, + ) + assert stage.outs[0].annot == Annotation(**annot) + assert (tmp_dir / "foo.dvc").parse() == M.dict(outs=[M.dict(**annot)]) + + # try to selectively update/overwrite some annotations + annot = {**annot, "type": "t2"} + stage = dvc.imp_url( + "remote://workspace/foo", + os.fspath(tmp_dir / "foo"), + no_exec=True, + type="t2", + ) + assert stage.outs[0].annot == Annotation(**annot) + assert (tmp_dir / "foo.dvc").parse() == M.dict(outs=[M.dict(**annot)]) diff --git a/tests/func/test_run_single_stage.py b/tests/func/test_run_single_stage.py index db775d051..7d939f5c1 100644 --- a/tests/func/test_run_single_stage.py +++ b/tests/func/test_run_single_stage.py @@ -972,6 +972,10 @@ def test_run_force_preserves_comments_and_meta(tmp_dir, dvc, run_copy): # comment preserved - path: bar desc: out desc + type: mytype + labels: + - label1 + - label2 meta: name: copy-foo-bar """ @@ -996,6 +1000,10 @@ def test_run_force_preserves_comments_and_meta(tmp_dir, dvc, run_copy): # comment preserved - path: bar desc: out desc + type: mytype + labels: + - label1 + - label2 md5: acbd18db4cc2f85cedef654fccc4a4d8 size: 3 meta: diff --git a/tests/func/test_stage.py b/tests/func/test_stage.py index 77a35488b..d9b72daf4 100644 --- a/tests/func/test_stage.py +++ b/tests/func/test_stage.py @@ -3,6 +3,7 @@ import tempfile import pytest +from dvc.annotations import Annotation from dvc.cli import main from dvc.dvcfile import SingleStageFile from dvc.fs import LocalFileSystem @@ -10,6 +11,8 @@ from dvc.output import Output from dvc.repo import Repo, lock_repo from dvc.stage import PipelineStage, Stage from dvc.stage.run import run_stage +from dvc.stage.utils import compute_md5 +from dvc.utils import dict_md5 from dvc.utils.serialize import dump_yaml, load_yaml from dvc.utils.strictyaml import YAMLValidationError from tests.basic_env import TestDvc @@ -164,40 +167,65 @@ def test_md5_ignores_comments(tmp_dir, dvc): assert not new_stage.changed_stage() -def test_meta_is_preserved(tmp_dir, dvc): - (stage,) = tmp_dir.dvc_gen("foo", "foo content") - - # Add meta to DVC-file - data = (tmp_dir / stage.path).parse() - data["meta"] = {"custom_key": 42} - (tmp_dir / stage.path).dump(data) - - # Loading and dumping to test that it works and meta is retained - dvcfile = SingleStageFile(dvc, stage.path) - new_stage = dvcfile.stage - dvcfile.dump(new_stage) - - new_data = (tmp_dir / stage.path).parse() - assert new_data["meta"] == data["meta"] +def test_md5_ignores_annotations(tmp_dir, dvc): + data = { + "desc": "stage desc", + "meta": {"key1": "value1", "key2": "value2"}, + "outs": [ + { + "md5": "d3b07384d113edec49eaa6238ad5ff00", + "size": 4, + "path": "foo", + "desc": "foo desc", + "type": "mytype", + "labels": ["get-started", "dataset-registry"], + } + ], + } + (tmp_dir / "foo.dvc").dump(data) + stage = dvc.stage.load_one("foo.dvc") + assert compute_md5(stage) == "1822617147b53ae6f9eb4b3c87c0b6f3" + assert ( + dict_md5( + { + "outs": [ + {"md5": "d3b07384d113edec49eaa6238ad5ff00", "path": "foo"} + ] + } + ) + == "1822617147b53ae6f9eb4b3c87c0b6f3" + ) -def test_desc_is_preserved(tmp_dir, dvc): - (stage,) = tmp_dir.dvc_gen("foo", "foo content") +def test_meta_desc_is_preserved(tmp_dir, dvc): + data = { + "desc": "stage desc", + "meta": {"key1": "value1", "key2": "value2"}, + "outs": [ + { + "md5": "d3b07384d113edec49eaa6238ad5ff00", + "size": 4, + "path": "foo", + "desc": "foo desc", + "type": "mytype", + "labels": ["get-started", "dataset-registry"], + } + ], + } + (tmp_dir / "foo.dvc").dump(data) + stage = dvc.stage.load_one("foo.dvc") - data = (tmp_dir / stage.path).parse() - stage_desc = "test stage description" - out_desc = "test out description" - data["desc"] = stage_desc - data["outs"][0]["desc"] = out_desc - (tmp_dir / stage.path).dump(data) + assert stage.meta == {"key1": "value1", "key2": "value2"} + assert stage.desc == "stage desc" + assert stage.outs[0].annot == Annotation( + desc="foo desc", + type="mytype", + labels=["get-started", "dataset-registry"], + ) - dvcfile = SingleStageFile(dvc, stage.path) - new_stage = dvcfile.stage - dvcfile.dump(new_stage) - - new_data = (tmp_dir / stage.path).parse() - assert new_data["desc"] == stage_desc - assert new_data["outs"][0]["desc"] == out_desc + # sanity check + stage.dump() + assert (tmp_dir / "foo.dvc").parse() == data def test_parent_repo_collect_stages(tmp_dir, scm, dvc): diff --git a/tests/unit/command/test_add.py b/tests/unit/command/test_add.py index 040b43b03..636f85981 100644 --- a/tests/unit/command/test_add.py +++ b/tests/unit/command/test_add.py @@ -37,6 +37,8 @@ def test_add(mocker, dvc): remote=None, to_remote=False, desc="stage description", + type=None, + labels=None, jobs=None, ) @@ -71,6 +73,8 @@ def test_add_to_remote(mocker): remote="remote", to_remote=True, desc=None, + type=None, + labels=None, jobs=None, ) diff --git a/tests/unit/command/test_annotations_flags.py b/tests/unit/command/test_annotations_flags.py new file mode 100644 index 000000000..e4fd0a783 --- /dev/null +++ b/tests/unit/command/test_annotations_flags.py @@ -0,0 +1,42 @@ +import pytest + +from dvc.cli import parse_args +from tests.utils.asserts import called_once_with_subset + + +@pytest.mark.parametrize( + "func, args", + [ + ("add", ("add", "model.pkl")), + ("imp", ("import", "https://github.com/org/repo", "model.pkl")), + ("imp_url", ("import-url", "https://dvc.org/data.xml", "data.xml")), + ], + ids=["add", "import", "import-url"], +) +def test_add_annotations(mocker, func, args): + mocker.patch("dvc.repo.Repo") + cli_args = parse_args( + [ + *args, + "--label", + "example-get-started", + "--label", + "model-registry", + "--type", + "model", + "--desc", + "description", + ] + ) + + _, *posargs = args + cmd = cli_args.func(cli_args) + + assert cmd.run() == 0 + assert called_once_with_subset( + getattr(cmd.repo, func), + posargs, + desc="description", + type="model", + labels=["example-get-started", "model-registry"], + ) diff --git a/tests/unit/command/test_imp.py b/tests/unit/command/test_imp.py index c8965e3a3..b078e1b15 100644 --- a/tests/unit/command/test_imp.py +++ b/tests/unit/command/test_imp.py @@ -36,6 +36,8 @@ def test_import(mocker): no_exec=False, no_download=False, desc="description", + type=None, + labels=None, jobs=3, ) @@ -72,6 +74,8 @@ def test_import_no_exec(mocker): no_exec=True, no_download=False, desc="description", + type=None, + labels=None, jobs=None, ) @@ -108,5 +112,7 @@ def test_import_no_download(mocker): no_exec=False, no_download=True, desc="description", + type=None, + labels=None, jobs=None, ) diff --git a/tests/unit/command/test_imp_url.py b/tests/unit/command/test_imp_url.py index 3f07f6a24..d723ccb21 100644 --- a/tests/unit/command/test_imp_url.py +++ b/tests/unit/command/test_imp_url.py @@ -37,6 +37,8 @@ def test_import_url(mocker): remote=None, to_remote=False, desc="description", + type=None, + labels=None, jobs=4, ) @@ -92,6 +94,8 @@ def test_import_url_no_exec_download_flags(mocker, flag, expected): to_remote=False, desc="description", jobs=None, + type=None, + labels=None, **expected ) @@ -125,6 +129,8 @@ def test_import_url_to_remote(mocker): remote="remote", to_remote=True, desc="description", + type=None, + labels=None, jobs=None, ) diff --git a/tests/unit/output/test_annotations.py b/tests/unit/output/test_annotations.py new file mode 100644 index 000000000..26b1de0d6 --- /dev/null +++ b/tests/unit/output/test_annotations.py @@ -0,0 +1,22 @@ +import pytest + +from dvc.annotations import Annotation + + +@pytest.mark.parametrize( + "kwargs", + [ + {"desc": "desc", "type": "type", "labels": ["label1", "label2"]}, + {"desc": "desc", "type": "type"}, + ], +) +def test_annotation_to_dict(kwargs): + annot = Annotation(**kwargs) + assert annot.to_dict() == kwargs + + +def test_annotation_update(): + annot = Annotation(desc="desc", labels=["label1", "label2"]) + annot.update(labels=["label"], type="type", new="new") + + assert vars(annot) == {"desc": "desc", "type": "type", "labels": ["label"]} diff --git a/tests/utils/asserts.py b/tests/utils/asserts.py index 084ee850b..579cf3124 100644 --- a/tests/utils/asserts.py +++ b/tests/utils/asserts.py @@ -11,7 +11,7 @@ def called_once_with_subset(m: Mock, *args: Any, **kwargs: Any) -> bool: m.assert_called_once() m_args, m_kwargs = m.call_args - expected_args = m_args + (ANY,) * (len(m_args) - len(args)) + expected_args = m_args + (ANY,) * (len(m_args) - len(args) - 1) expected_kwargs = {k: kwargs.get(k, ANY) for k in m_kwargs} m.assert_called_with(*expected_args, **expected_kwargs)