From c73764156eeea688424bbee669b4279a8bb89b96 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Saugat=20Pachhai=20=28=E0=A4=B8=E0=A5=8C=E0=A4=97=E0=A4=BE?= =?UTF-8?q?=E0=A4=A4=29?= Date: Tue, 6 Sep 2022 17:09:21 +0545 Subject: [PATCH] metadata: support custom metadata in .dvc files This adds support for custom metadata in outputs under `meta` key. Also adds `--meta` flag to `add`/`import`/`import-url` that supports `key=value` values. --- dvc/annotations.py | 5 ++++- dvc/cli/actions.py | 11 +++++++++++ dvc/commands/add.py | 9 +++++++++ dvc/commands/imp.py | 1 + dvc/commands/imp_url.py | 1 + dvc/output.py | 5 ++++- dvc/repo/imp_url.py | 3 ++- tests/func/test_add.py | 7 ++++++- tests/func/test_commit.py | 6 ++++++ tests/func/test_import.py | 7 ++++++- tests/func/test_import_url.py | 11 ++++++++++- tests/func/test_run_single_stage.py | 4 ++++ tests/func/test_stage.py | 3 +++ tests/unit/command/test_add.py | 2 ++ tests/unit/command/test_annotations_flags.py | 5 +++++ tests/unit/command/test_imp.py | 3 +++ tests/unit/command/test_imp_url.py | 3 +++ tests/unit/output/test_annotations.py | 13 ++++++++++--- 18 files changed, 90 insertions(+), 9 deletions(-) create mode 100644 dvc/cli/actions.py diff --git a/dvc/annotations.py b/dvc/annotations.py index eb3b856a0..05b02fb89 100644 --- a/dvc/annotations.py +++ b/dvc/annotations.py @@ -1,5 +1,5 @@ from dataclasses import asdict, dataclass, field, fields -from typing import ClassVar, Dict, List, Optional +from typing import Any, ClassVar, Dict, List, Optional from funcy import compact @@ -9,10 +9,12 @@ class Annotation: PARAM_DESC: ClassVar[str] = "desc" PARAM_TYPE: ClassVar[str] = "type" PARAM_LABELS: ClassVar[str] = "labels" + PARAM_META: ClassVar[str] = "meta" desc: Optional[str] = None type: Optional[str] = None labels: List[str] = field(default_factory=list) + meta: Dict[str, Any] = field(default_factory=dict) def update(self, **kwargs) -> "Annotation": for attr, value in kwargs.items(): @@ -29,4 +31,5 @@ ANNOTATION_SCHEMA = { Annotation.PARAM_DESC: str, Annotation.PARAM_TYPE: str, Annotation.PARAM_LABELS: [str], + Annotation.PARAM_META: object, } diff --git a/dvc/cli/actions.py b/dvc/cli/actions.py new file mode 100644 index 000000000..279b53680 --- /dev/null +++ b/dvc/cli/actions.py @@ -0,0 +1,11 @@ +from argparse import _AppendAction + + +class KeyValueArgs(_AppendAction): + def __call__(self, parser, namespace, values, option_string=None): + items = getattr(namespace, self.dest) or {} + for value in filter(bool, values): + key, _, value = value.partition("=") + items[key.strip()] = value + + setattr(namespace, self.dest, items) diff --git a/dvc/commands/add.py b/dvc/commands/add.py index 091519bd7..8c9aeb842 100644 --- a/dvc/commands/add.py +++ b/dvc/commands/add.py @@ -2,6 +2,7 @@ import argparse import logging from dvc.cli import completion +from dvc.cli.actions import KeyValueArgs from dvc.cli.command import CmdBase from dvc.cli.utils import append_doc_link @@ -18,6 +19,13 @@ def _add_annotating_args(parser: argparse.ArgumentParser) -> None: "This doesn't affect any DVC operations." ), ) + parser.add_argument( + "--meta", + metavar="key=value", + nargs=1, + action=KeyValueArgs, + help="Custom metadata to add to the data", + ) parser.add_argument( "--label", dest="labels", @@ -56,6 +64,7 @@ class CmdAdd(CmdBase): out=self.args.out, type=self.args.type, labels=self.args.labels, + meta=self.args.meta, remote=self.args.remote, to_remote=self.args.to_remote, jobs=self.args.jobs, diff --git a/dvc/commands/imp.py b/dvc/commands/imp.py index bfa75445c..f11b5bb83 100644 --- a/dvc/commands/imp.py +++ b/dvc/commands/imp.py @@ -23,6 +23,7 @@ class CmdImport(CmdBase): desc=self.args.desc, type=self.args.type, labels=self.args.labels, + meta=self.args.meta, jobs=self.args.jobs, ) except DvcException: diff --git a/dvc/commands/imp_url.py b/dvc/commands/imp_url.py index baa6e681a..370f6c557 100644 --- a/dvc/commands/imp_url.py +++ b/dvc/commands/imp_url.py @@ -23,6 +23,7 @@ class CmdImportUrl(CmdBase): desc=self.args.desc, type=self.args.type, labels=self.args.labels, + meta=self.args.meta, jobs=self.args.jobs, ) except DvcException: diff --git a/dvc/output.py b/dvc/output.py index dd2dd929d..ff719ae32 100644 --- a/dvc/output.py +++ b/dvc/output.py @@ -287,10 +287,13 @@ class Output: desc=None, type=None, # pylint: disable=redefined-builtin labels=None, + meta=None, remote=None, repo=None, ): - self.annot = Annotation(desc=desc, type=type, labels=labels or []) + self.annot = Annotation( + desc=desc, type=type, labels=labels or [], meta=meta 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 diff --git a/dvc/repo/imp_url.py b/dvc/repo/imp_url.py index ae1d1a0f1..d3c1e2390 100644 --- a/dvc/repo/imp_url.py +++ b/dvc/repo/imp_url.py @@ -24,6 +24,7 @@ def imp_url( desc=None, type=None, # pylint: disable=redefined-builtin labels=None, + meta=None, jobs=None, ): from dvc.dvcfile import Dvcfile @@ -64,7 +65,7 @@ def imp_url( restore_fields(stage) out_obj = stage.outs[0] - out_obj.annot.update(desc=desc, type=type, labels=labels) + out_obj.annot.update(desc=desc, type=type, labels=labels, meta=meta) dvcfile = Dvcfile(self, stage.path) dvcfile.remove() diff --git a/tests/func/test_add.py b/tests/func/test_add.py index 863ba206e..0c628fd29 100644 --- a/tests/func/test_add.py +++ b/tests/func/test_add.py @@ -1175,7 +1175,12 @@ def test_add_ignore_duplicated_targets(tmp_dir, dvc, capsys): def test_add_with_annotations(M, tmp_dir, dvc): tmp_dir.gen("foo", "foo") - annot = {"desc": "foo desc", "labels": ["l1", "l2"], "type": "t1"} + annot = { + "desc": "foo desc", + "labels": ["l1", "l2"], + "type": "t1", + "meta": {"key": "value"}, + } (stage,) = dvc.add("foo", **annot) 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 b6575be6d..e24330e1d 100644 --- a/tests/func/test_commit.py +++ b/tests/func/test_commit.py @@ -51,6 +51,9 @@ def test_commit_preserve_fields(tmp_dir, dvc): labels: - label1 - label2 + meta: + key1: value1 + key2: value2 remote: testremote meta: some metadata """ @@ -69,6 +72,9 @@ def test_commit_preserve_fields(tmp_dir, dvc): labels: - label1 - label2 + meta: + key1: value1 + key2: value2 remote: testremote md5: acbd18db4cc2f85cedef654fccc4a4d8 size: 3 diff --git a/tests/func/test_import.py b/tests/func/test_import.py index 5f01d24d1..94fc343f9 100644 --- a/tests/func/test_import.py +++ b/tests/func/test_import.py @@ -634,7 +634,12 @@ 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"} + annot = { + "desc": "foo desc", + "labels": ["l1", "l2"], + "type": "t1", + "meta": {"key": "value"}, + } 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)]) diff --git a/tests/func/test_import_url.py b/tests/func/test_import_url.py index 55e13a2e9..6e34de8b3 100644 --- a/tests/func/test_import_url.py +++ b/tests/func/test_import_url.py @@ -147,6 +147,8 @@ def test_import_url_preserve_fields(tmp_dir, dvc): labels: - label1 - label2 + meta: + key: value meta: some metadata """ ) @@ -169,6 +171,8 @@ def test_import_url_preserve_fields(tmp_dir, dvc): labels: - label1 - label2 + meta: + key: value md5: acbd18db4cc2f85cedef654fccc4a4d8 size: 3 meta: some metadata @@ -301,7 +305,12 @@ def test_import_url_no_download(tmp_dir, dvc, local_workspace): 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"} + annot = { + "desc": "foo desc", + "labels": ["l1", "l2"], + "type": "t1", + "meta": {"key": "value"}, + } stage = dvc.imp_url( "remote://workspace/foo", os.fspath(tmp_dir / "foo"), diff --git a/tests/func/test_run_single_stage.py b/tests/func/test_run_single_stage.py index 7d939f5c1..193122ad2 100644 --- a/tests/func/test_run_single_stage.py +++ b/tests/func/test_run_single_stage.py @@ -976,6 +976,8 @@ def test_run_force_preserves_comments_and_meta(tmp_dir, dvc, run_copy): labels: - label1 - label2 + meta: + key: value meta: name: copy-foo-bar """ @@ -1004,6 +1006,8 @@ def test_run_force_preserves_comments_and_meta(tmp_dir, dvc, run_copy): labels: - label1 - label2 + meta: + key: value md5: acbd18db4cc2f85cedef654fccc4a4d8 size: 3 meta: diff --git a/tests/func/test_stage.py b/tests/func/test_stage.py index d9b72daf4..82950102b 100644 --- a/tests/func/test_stage.py +++ b/tests/func/test_stage.py @@ -179,6 +179,7 @@ def test_md5_ignores_annotations(tmp_dir, dvc): "desc": "foo desc", "type": "mytype", "labels": ["get-started", "dataset-registry"], + "meta": {"key1": "value1"}, } ], } @@ -209,6 +210,7 @@ def test_meta_desc_is_preserved(tmp_dir, dvc): "desc": "foo desc", "type": "mytype", "labels": ["get-started", "dataset-registry"], + "meta": {"key": "value"}, } ], } @@ -221,6 +223,7 @@ def test_meta_desc_is_preserved(tmp_dir, dvc): desc="foo desc", type="mytype", labels=["get-started", "dataset-registry"], + meta={"key": "value"}, ) # sanity check diff --git a/tests/unit/command/test_add.py b/tests/unit/command/test_add.py index 636f85981..69581587d 100644 --- a/tests/unit/command/test_add.py +++ b/tests/unit/command/test_add.py @@ -39,6 +39,7 @@ def test_add(mocker, dvc): desc="stage description", type=None, labels=None, + meta=None, jobs=None, ) @@ -75,6 +76,7 @@ def test_add_to_remote(mocker): desc=None, type=None, labels=None, + meta=None, jobs=None, ) diff --git a/tests/unit/command/test_annotations_flags.py b/tests/unit/command/test_annotations_flags.py index e4fd0a783..027d9a104 100644 --- a/tests/unit/command/test_annotations_flags.py +++ b/tests/unit/command/test_annotations_flags.py @@ -26,6 +26,10 @@ def test_add_annotations(mocker, func, args): "model", "--desc", "description", + "--meta", + "key1=value1", + "--meta", + "key2=value2", ] ) @@ -39,4 +43,5 @@ def test_add_annotations(mocker, func, args): desc="description", type="model", labels=["example-get-started", "model-registry"], + meta={"key1": "value1", "key2": "value2"}, ) diff --git a/tests/unit/command/test_imp.py b/tests/unit/command/test_imp.py index b078e1b15..11a286edd 100644 --- a/tests/unit/command/test_imp.py +++ b/tests/unit/command/test_imp.py @@ -38,6 +38,7 @@ def test_import(mocker): desc="description", type=None, labels=None, + meta=None, jobs=3, ) @@ -76,6 +77,7 @@ def test_import_no_exec(mocker): desc="description", type=None, labels=None, + meta=None, jobs=None, ) @@ -114,5 +116,6 @@ def test_import_no_download(mocker): desc="description", type=None, labels=None, + meta=None, jobs=None, ) diff --git a/tests/unit/command/test_imp_url.py b/tests/unit/command/test_imp_url.py index d723ccb21..6249d7527 100644 --- a/tests/unit/command/test_imp_url.py +++ b/tests/unit/command/test_imp_url.py @@ -39,6 +39,7 @@ def test_import_url(mocker): desc="description", type=None, labels=None, + meta=None, jobs=4, ) @@ -96,6 +97,7 @@ def test_import_url_no_exec_download_flags(mocker, flag, expected): jobs=None, type=None, labels=None, + meta=None, **expected ) @@ -131,6 +133,7 @@ def test_import_url_to_remote(mocker): desc="description", type=None, labels=None, + meta=None, jobs=None, ) diff --git a/tests/unit/output/test_annotations.py b/tests/unit/output/test_annotations.py index 26b1de0d6..1486a2d3c 100644 --- a/tests/unit/output/test_annotations.py +++ b/tests/unit/output/test_annotations.py @@ -7,7 +7,7 @@ from dvc.annotations import Annotation "kwargs", [ {"desc": "desc", "type": "type", "labels": ["label1", "label2"]}, - {"desc": "desc", "type": "type"}, + {"desc": "desc", "type": "type", "meta": {"key": "value"}}, ], ) def test_annotation_to_dict(kwargs): @@ -17,6 +17,13 @@ def test_annotation_to_dict(kwargs): def test_annotation_update(): annot = Annotation(desc="desc", labels=["label1", "label2"]) - annot.update(labels=["label"], type="type", new="new") + annot.update( + labels=["label"], type="type", new="new", meta={"key": "value"} + ) - assert vars(annot) == {"desc": "desc", "type": "type", "labels": ["label"]} + assert vars(annot) == { + "desc": "desc", + "type": "type", + "labels": ["label"], + "meta": {"key": "value"}, + }