Stone_Red/dvc
1
0
Fork 0
mirror of https://github.com/Stone-Red-Code/dvc.git synced 2026-11-02 07:18:25 +01:00
Code Issues Packages Projects Releases Wiki Activity

add metadata fields: label, type to data (#8232)

Browse Source 
gto like metadata fields: labels and type
This commit is contained in:
skshetry
2022-09-08 16:33:47 +05:45
committed by GitHub
parent 504485ed6e
commit df60df5242
24 changed files with 344 additions and 93 deletions
Download Patch File Download Diff File Expand all files Collapse all files
dvc/annotations.py
+32
View File
@@ -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],
}
dvc/commands/add.py
+30 -9
View File
@@ -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="<text>",
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="<str>",
help="Comma separated list of labels for the data (optional).",
)
parser.add_argument(
"--type",
type=str,
metavar="<str>",
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="<number>",
)
parser.add_argument(
"--desc",
type=str,
metavar="<text>",
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
dvc/commands/imp.py
+6 -9
View File
@@ -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="<text>",
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="<number>",
)
_add_annotating_args(import_parser)
import_parser.set_defaults(func=CmdImport)
dvc/commands/imp_url.py
+5 -9
View File
@@ -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="<number>",
)
import_parser.add_argument(
"--desc",
type=str,
metavar="<text>",
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)
dvc/output.py
+19 -15
View File
@@ -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,
}
dvc/repo/add.py
+3 -2
View File
@@ -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
dvc/repo/imp_url.py
+4 -3
View File
@@ -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()
dvc/schema.py
+2 -1
View File
@@ -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,
}
}
dvc/stage/__init__.py
+3 -6
View File
@@ -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):
dvc/stage/serialize.py
+4 -4
View File
@@ -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:
dvc/stage/utils.py
+2 -1
View File
@@ -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,
tests/func/test_add.py
+24
View File
@@ -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)])
tests/func/test_commit.py
+8
View File
@@ -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
tests/func/test_dvcfile.py
+2 -2
View File
@@ -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()
tests/func/test_import.py
+19
View File
@@ -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)])
tests/func/test_import_url.py
+34 -1
View File
@@ -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)])
tests/func/test_run_single_stage.py
+8
View File
@@ -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:
tests/func/test_stage.py
+58 -30
View File
@@ -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):
tests/unit/command/test_add.py
+4
View File
@@ -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,
)
tests/unit/command/test_annotations_flags.py
+42
View File
@@ -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"],
)
tests/unit/command/test_imp.py
+6
View File
@@ -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,
)
tests/unit/command/test_imp_url.py
+6
View File
@@ -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,
)
tests/unit/output/test_annotations.py
+22
View File
@@ -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"]}
tests/utils/asserts.py
+1 -1
View File
@@ -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)