Add API to add artifacts (#9345)

* api call to add/write artifact

* return dvcyaml path on add

* create dir for dvc.yaml if needed

* create dir for dvc.yaml if needed - fix

* allow nested artifact names

* make dvc get download artifacts

* Apply suggestions from code review

Co-authored-by: Oded Messer <[email protected]>

* use feedback

* remove 'dvc get' from this PR, keep only adding artifacts

* allow first char of artifact name to be a number

* add note about constants

* fix parsing artifact name

* remove 'parse_name' and write tests for 'add'

* Apply suggestions from code review

Co-authored-by: Oded Messer <[email protected]>

* updates from comments

* fix for studio

* fix again

* making check_name_format to raise an exception

* Update dvc/repo/artifacts.py

---------

Co-authored-by: Oded Messer <[email protected]>

dvc/annotations.py

@@ -28,9 +28,21 @@ class Annotation:
 
 
 @dataclass
-class Artifact(Annotation):
+class Artifact:
     PARAM_PATH: ClassVar[str] = "path"
-    path: Optional[str] = None
+    PARAM_DESC: ClassVar[str] = "desc"
+    PARAM_TYPE: ClassVar[str] = "type"
+    PARAM_LABELS: ClassVar[str] = "labels"
+    PARAM_META: ClassVar[str] = "meta"
+
+    path: str
+    desc: Optional[str] = None
+    type: Optional[str] = None  # noqa: A003
+    labels: List[str] = field(default_factory=list)
+    meta: Dict[str, Any] = field(default_factory=dict)
+
+    def to_dict(self) -> Dict[str, str]:
+        return compact(asdict(self))
 
 
 ANNOTATION_FIELDS = [field.name for field in fields(Annotation)]

dvc/repo/artifacts.py

@@ -1,19 +1,26 @@
 import logging
-import os
 import re
-from typing import TYPE_CHECKING, Dict
+from pathlib import Path
+from typing import Dict, Optional
 
 from dvc.annotations import Artifact
-from dvc.dvcfile import FileMixin
+from dvc.dvcfile import PROJECT_FILE
+from dvc.exceptions import InvalidArgumentError
+from dvc.repo import Repo
 from dvc.utils import relpath
-
-if TYPE_CHECKING:
-    from dvc.repo import Repo
+from dvc.utils.serialize import modify_yaml
 
 logger = logging.getLogger(__name__)
 
 
-NAME_RE = re.compile(r"^[a-z]([a-z0-9-]*[a-z0-9])?$")
+# Constants are taken from GTO.
+# When we make it a dependency, we can import them instead
+SEPARATOR_IN_NAME = ":"
+DIRNAME = r"[a-z0-9-_./]+"
+NAME = r"[a-z0-9]([a-z0-9-/]*[a-z0-9])?"
+NAME_RE = re.compile(f"^{NAME}$")
+FULLNAME = f"((?P<dirname>{DIRNAME}){SEPARATOR_IN_NAME})?(?P<name>{NAME})"
+FULLNAME_RE = re.compile(f"^{FULLNAME}$")
 
 
 def name_is_compatible(name: str) -> bool:
@@ -22,22 +29,23 @@ def name_is_compatible(name: str) -> bool:
 
 def check_name_format(name: str) -> None:
     if not name_is_compatible(name):
-        logger.warning(
-            "Can't use '%s' as artifact name (ID)."
+        raise InvalidArgumentError(
+            f"Can't use '{name}' as artifact name (ID)."
             " You can use letters and numbers, and use '-' as separator"
-            " (but not at the start or end). The first character must be a letter.",
-            name,
+            " (but not at the start or end)."
         )
 
 
-class ArtifactsFile(FileMixin):
-    from dvc.schema import SINGLE_ARTIFACT_SCHEMA as SCHEMA
-
-    def dump(self, stage, **kwargs):
-        raise NotImplementedError
-
-    def merge(self, ancestor, other, allowed=None):
-        raise NotImplementedError
+def check_for_nested_dvc_repo(dvcfile: Path):
+    if dvcfile.is_absolute():
+        raise InvalidArgumentError("Use relative path to dvc.yaml.")
+    path = dvcfile.parent
+    while path.name:
+        if (path / Repo.DVC_DIR).is_dir():
+            raise InvalidArgumentError(
+                f"Nested DVC repos like {path} are not supported."
+            )
+        path = path.parent
 
 
 class Artifacts:
@@ -50,18 +58,30 @@ class Artifacts:
             dvcfile,
             dvcfile_artifacts,
         ) in self.repo.index._artifacts.items():  # pylint: disable=protected-access
-            # read the artifacts.yaml file if needed
-            if isinstance(dvcfile_artifacts, str):
-                dvcfile_artifacts = ArtifactsFile(
-                    self.repo,
-                    os.path.join(os.path.dirname(dvcfile), dvcfile_artifacts),
-                    verify=False,
-                ).load()
-            if not dvcfile_artifacts:
-                continue
             dvcyaml = relpath(dvcfile, self.repo.root_dir)
             artifacts[dvcyaml] = {}
             for name, value in dvcfile_artifacts.items():
-                check_name_format(name)
+                try:
+                    check_name_format(name)
+                except InvalidArgumentError as e:
+                    logger.warning(e.msg)
                 artifacts[dvcyaml][name] = Artifact(**value)
         return artifacts
+
+    def add(self, name: str, artifact: Artifact, dvcfile: Optional[str] = None):
+        with self.repo.scm_context(quiet=True):
+            check_name_format(name)
+            dvcyaml = Path(dvcfile or PROJECT_FILE)
+            check_for_nested_dvc_repo(
+                dvcyaml.relative_to(self.repo.root_dir)
+                if dvcyaml.is_absolute()
+                else dvcyaml
+            )
+
+            with modify_yaml(dvcyaml) as data:
+                artifacts = data.setdefault("artifacts", {})
+                artifacts.update({name: artifact.to_dict()})
+
+            self.repo.scm_context.track_file(dvcfile)
+
+        return artifacts.get(name)

dvc/schema.py

@@ -116,7 +116,6 @@ PLOT_DEFINITION = {
 SINGLE_PLOT_SCHEMA = {str: Any(PLOT_DEFINITION, None)}
 ARTIFACTS = "artifacts"
 SINGLE_ARTIFACT_SCHEMA = Schema({str: ARTIFACT_SCHEMA})
-ARTIFACTS_SCHEMA = Any(str, SINGLE_ARTIFACT_SCHEMA)
 FOREACH_IN = {
     Required(FOREACH_KWD): Any(dict, list, str),
     Required(DO_KWD): STAGE_DEFINITION,
@@ -130,7 +129,7 @@ MULTI_STAGE_SCHEMA = {
     VARS_KWD: VARS_SCHEMA,
     StageParams.PARAM_PARAMS: [str],
     StageParams.PARAM_METRICS: [str],
-    ARTIFACTS: ARTIFACTS_SCHEMA,
+    ARTIFACTS: SINGLE_ARTIFACT_SCHEMA,
 }
 
 COMPILED_SINGLE_STAGE_SCHEMA = Schema(SINGLE_STAGE_SCHEMA)

tests/func/artifacts/test_artifacts.py

@@ -3,6 +3,7 @@ import os
 import pytest
 
 from dvc.annotations import Artifact
+from dvc.exceptions import InvalidArgumentError
 from dvc.repo.artifacts import name_is_compatible
 from dvc.utils.strictyaml import YAMLSyntaxError, YAMLValidationError
 
@@ -21,7 +22,7 @@ dvcyaml = {
 }
 
 
-def test_reading_artifacts_subdir(tmp_dir, dvc):
+def test_artifacts_read_subdir(tmp_dir, dvc):
     (tmp_dir / "dvc.yaml").dump(dvcyaml)
 
     subdir = tmp_dir / "subdir"
@@ -38,6 +39,55 @@ def test_reading_artifacts_subdir(tmp_dir, dvc):
     }
 
 
+def test_artifacts_add_subdir(tmp_dir, dvc):
+    subdir = tmp_dir / "subdir"
+    subdir.mkdir()
+
+    (subdir / "dvc.yaml").dump(dvcyaml)
+
+    new_art = Artifact(path="path")
+    tmp_dir.dvc.artifacts.add("new", new_art, dvcfile="subdir/dvc.yaml")
+
+    artifacts = {
+        name: Artifact(**values) for name, values in dvcyaml["artifacts"].items()
+    }
+    artifacts["new"] = new_art
+    assert tmp_dir.dvc.artifacts.read() == {
+        f"subdir{os.path.sep}dvc.yaml": artifacts,
+    }
+
+
+def test_artifacts_add_abspath(tmp_dir, dvc):
+    subdir = tmp_dir / "subdir"
+    subdir.mkdir()
+
+    new_art = Artifact(path="path")
+    tmp_dir.dvc.artifacts.add(
+        "new", new_art, dvcfile=os.path.abspath("subdir/dvc.yaml")
+    )
+
+    assert tmp_dir.dvc.artifacts.read() == {
+        f"subdir{os.path.sep}dvc.yaml": {"new": new_art},
+    }
+
+
+def test_artifacts_add_fails_on_dvc_subrepo(tmp_dir, dvc):
+    # adding artifact to the DVC subrepo from the parent DVC repo
+    # shouldn't work
+    subdir = tmp_dir / "subdir"
+    (subdir / ".dvc").mkdir(parents=True)
+
+    with pytest.raises(InvalidArgumentError):
+        tmp_dir.dvc.artifacts.add(
+            "failing", Artifact(path="path"), dvcfile="subdir/dvc.yaml"
+        )
+
+    with pytest.raises(InvalidArgumentError):
+        tmp_dir.dvc.artifacts.add(
+            "failing", Artifact(path="path"), dvcfile="subdir/dvclive/dvc.yaml"
+        )
+
+
 bad_dvcyaml_extra_field = {
     "artifacts": {
         "lol": {"kek": "cheburek", "path": "lol"},
@@ -71,7 +121,7 @@ artifacts:
 """
 
 
-def test_broken_dvcyaml_id_duplication(tmp_dir, dvc):
+def test_artifacts_read_fails_on_id_duplication(tmp_dir, dvc):
     with open(tmp_dir / "dvc.yaml", "w") as f:
         f.write(bad_dvcyaml_id_duplication)
 
@@ -79,32 +129,19 @@ def test_broken_dvcyaml_id_duplication(tmp_dir, dvc):
         tmp_dir.dvc.artifacts.read()
 
 
-dvcyaml_redirecting = {"artifacts": "artifacts.yaml"}
-
-
-def test_read_artifacts_yaml(tmp_dir, dvc):
-    (tmp_dir / "dvc.yaml").dump(dvcyaml_redirecting)
-    (tmp_dir / "artifacts.yaml").dump(dvcyaml["artifacts"])
-
-    artifacts = {
-        name: Artifact(**values) for name, values in dvcyaml["artifacts"].items()
-    }
-    assert tmp_dir.dvc.artifacts.read() == {
-        "dvc.yaml": artifacts,
-    }
-
-
 @pytest.mark.parametrize(
     "name",
     [
+        "1",
         "m",
         "nn",
         "m1",
+        "1nn",
         "model-prod",
         "model-prod-v1",
     ],
 )
-def test_check_name_is_valid(name):
+def test_name_is_compatible(name):
     assert name_is_compatible(name)
 
 
@@ -112,10 +149,8 @@ def test_check_name_is_valid(name):
     "name",
     [
         "",
-        "1",
         "m/",
         "/m",
-        "1nn",
         "###",
         "@@@",
         "a model",
@@ -125,8 +160,7 @@ def test_check_name_is_valid(name):
         "model@1",
         "model#1",
         "@namespace/model",
-        "namespace/model",
     ],
 )
-def test_check_name_is_invalid(name):
+def test_name_is_compatible_fails(name):
     assert not name_is_compatible(name)