dvc: introduce dvc tag

Fixes #1766

Signed-off-by: Ruslan Kuprieiev <[email protected]>

dvc/cli.py

@@ -31,6 +31,7 @@ import dvc.command.lock as lock
 import dvc.command.pipeline as pipeline
 import dvc.command.daemon as daemon
 import dvc.command.commit as commit
+import dvc.command.tag as tag
 from dvc.exceptions import DvcParserError
 from dvc import VERSION
 
@@ -60,6 +61,7 @@ COMMANDS = [
     pipeline,
     daemon,
     commit,
+    tag,
 ]
 
 

dvc/command/tag.py

@@ -0,0 +1,151 @@
+import yaml
+
+import dvc.logger as logger
+from dvc.exceptions import DvcException
+from dvc.command.base import CmdBase, fix_subparsers
+
+
+class CmdTagAdd(CmdBase):
+    def run(self):
+        for target in self.args.targets:
+            try:
+                self.repo.tag.add(
+                    self.args.tag,
+                    target=target,
+                    with_deps=self.args.with_deps,
+                    recursive=self.args.recursive,
+                )
+            except DvcException:
+                logger.error("failed to add tag")
+                return 1
+        return 0
+
+
+class CmdTagRemove(CmdBase):
+    def run(self):
+        for target in self.args.targets:
+            try:
+                self.repo.tag.remove(
+                    self.args.tag,
+                    target=target,
+                    with_deps=self.args.with_deps,
+                    recursive=self.args.recursive,
+                )
+            except DvcException:
+                logger.error("failed to remove tag")
+                return 1
+        return 0
+
+
+class CmdTagList(CmdBase):
+    def run(self):
+        for target in self.args.targets:
+            try:
+                tags = self.repo.tag.list(
+                    target,
+                    with_deps=self.args.with_deps,
+                    recursive=self.args.recursive,
+                )
+                if tags:
+                    logger.info(yaml.dump(tags, default_flow_style=False))
+            except DvcException:
+                logger.error("failed list tags")
+                return 1
+        return 0
+
+
+def add_parser(subparsers, parent_parser):
+    TAG_HELP = "A set of commands to manage dvc tags."
+    tag_parser = subparsers.add_parser(
+        "tag",
+        parents=[parent_parser],
+        description=append_doc_link(TAG_HELP, "tag"),
+        add_help=TAG_HELP,
+    )
+
+    tag_subparsers = tag_parser.add_subparsers(
+        dest="cmd",
+        help="Use dvc tag CMD --help to display command-specific help.",
+    )
+
+    fix_subparsers(tag_subparsers)
+
+    TAG_ADD_HELP = "Add dvc tag."
+    tag_add_parser = tag_subparsers.add_parser(
+        "add",
+        parents=[parent_parser],
+        description=TAG_ADD_HELP,
+        help=TAG_ADD_HELP,
+    )
+    tag_add_parser.add_argument("tag", help="Dvc tag.")
+    tag_add_parser.add_argument(
+        "targets", nargs="*", default=[None], help="Dvc files."
+    )
+    tag_add_parser.add_argument(
+        "-d",
+        "--with-deps",
+        action="store_true",
+        default=False,
+        help="Add tag for all dependencies of the specified dvc file.",
+    )
+    tag_add_parser.add_argument(
+        "-R",
+        "--recursive",
+        action="store_true",
+        default=False,
+        help="Add tag for subdirectories of the specified directory.",
+    )
+    tag_add_parser.set_defaults(func=CmdTagAdd)
+
+    TAG_REMOVE_HELP = "Remove dvc tag."
+    tag_remove_parser = tag_subparsers.add_parser(
+        "remove",
+        parents=[parent_parser],
+        description=TAG_REMOVE_HELP,
+        help=TAG_REMOVE_HELP,
+    )
+    tag_remove_parser.add_argument("tag", help="Dvc tag.")
+    tag_remove_parser.add_argument(
+        "targets", nargs="*", default=[None], help="Dvc files."
+    )
+    tag_remove_parser.add_argument(
+        "-d",
+        "--with-deps",
+        action="store_true",
+        default=False,
+        help="Remove tag for all dependencies of the specified dvc file.",
+    )
+    tag_remove_parser.add_argument(
+        "-R",
+        "--recursive",
+        action="store_true",
+        default=False,
+        help="Remove tag for subdirectories of the specified directory.",
+    )
+    tag_remove_parser.set_defaults(func=CmdTagRemove)
+
+    TAG_LIST_HELP = "List dvc tags."
+    tag_list_parser = tag_subparsers.add_parser(
+        "list",
+        parents=[parent_parser],
+        description=TAG_LIST_HELP,
+        help=TAG_LIST_HELP,
+    )
+    tag_list_parser.add_argument(
+        "targets", nargs="*", default=[None], help="Dvc files."
+    )
+    tag_list_parser.add_argument(
+        "-d",
+        "--with-deps",
+        action="store_true",
+        default=False,
+        help="List tags for all dependencies of the specified dvc file.",
+    )
+    tag_list_parser.add_argument(
+        "-R",
+        "--recursive",
+        action="store_true",
+        default=False,
+        help="List tags for subdirectories of the specified directory.",
+    )
+    tag_list_parser.set_defaults(func=CmdTagList)

dvc/output/__init__.py

@@ -33,26 +33,31 @@ OUTS_MAP = {
     "local": OutputLOCAL,
 }
 
-SCHEMA = {
-    OutputBase.PARAM_PATH: str,
-    # NOTE: currently there are only 3 possible checksum names:
-    #
-    #    1) md5 (LOCAL, SSH, GS);
-    #    2) etag (S3);
-    #    3) checksum (HDFS);
-    #
-    # so when a few types of outputs share the same name, we only need
-    # specify it once.
+# NOTE: currently there are only 3 possible checksum names:
+#
+#    1) md5 (LOCAL, SSH, GS);
+#    2) etag (S3);
+#    3) checksum (HDFS);
+#
+# so when a few types of outputs share the same name, we only need
+# specify it once.
+CHECKSUM_SCHEMA = {
     schema.Optional(RemoteLOCAL.PARAM_CHECKSUM): schema.Or(str, None),
     schema.Optional(RemoteS3.PARAM_CHECKSUM): schema.Or(str, None),
     schema.Optional(RemoteHDFS.PARAM_CHECKSUM): schema.Or(str, None),
-    schema.Optional(OutputBase.PARAM_CACHE): bool,
-    schema.Optional(OutputBase.PARAM_METRIC): OutputBase.METRIC_SCHEMA,
-    schema.Optional(OutputBase.PARAM_PERSIST): bool,
 }
 
+TAGS_SCHEMA = {schema.Optional(str): CHECKSUM_SCHEMA}
 
-def _get(stage, p, info, cache, metric, persist):
+SCHEMA = CHECKSUM_SCHEMA.copy()
+SCHEMA[OutputBase.PARAM_PATH] = str
+SCHEMA[schema.Optional(OutputBase.PARAM_CACHE)] = bool
+SCHEMA[schema.Optional(OutputBase.PARAM_METRIC)] = OutputBase.METRIC_SCHEMA
+SCHEMA[schema.Optional(OutputBase.PARAM_TAGS)] = TAGS_SCHEMA
+SCHEMA[schema.Optional(OutputBase.PARAM_PERSIST)] = bool
+
+
+def _get(stage, p, info, cache, metric, persist=False, tags=None):
     parsed = urlparse(p)
     if parsed.scheme == "remote":
         name = Config.SECTION_REMOTE_FMT.format(parsed.netloc)
@@ -66,11 +71,21 @@ def _get(stage, p, info, cache, metric, persist):
             remote=remote,
             metric=metric,
             persist=persist,
+            tags=tags,
         )
 
     for o in OUTS:
         if o.supported(p):
-            return o(stage, p, info, cache=cache, remote=None, metric=metric)
+            return o(
+                stage,
+                p,
+                info,
+                cache=cache,
+                remote=None,
+                metric=metric,
+                persist=persist,
+                tags=tags,
+            )
     return OutputLOCAL(
         stage,
         p,
@@ -79,6 +94,7 @@ def _get(stage, p, info, cache, metric, persist):
         remote=None,
         metric=metric,
         persist=persist,
+        tags=tags,
     )
 
 
@@ -89,8 +105,17 @@ def loadd_from(stage, d_list):
         cache = d.pop(OutputBase.PARAM_CACHE, True)
         metric = d.pop(OutputBase.PARAM_METRIC, False)
         persist = d.pop(OutputBase.PARAM_PERSIST, False)
+        tags = d.pop(OutputBase.PARAM_TAGS, None)
         ret.append(
-            _get(stage, p, info=d, cache=cache, metric=metric, persist=persist)
+            _get(
+                stage,
+                p,
+                info=d,
+                cache=cache,
+                metric=metric,
+                persist=persist,
+                tags=tags,
+            )
         )
     return ret
 

dvc/output/base.py

@@ -47,6 +47,8 @@ class OutputBase(object):
         },
     )
 
+    PARAM_TAGS = "tags"
+
     DoesNotExistError = OutputDoesNotExistError
     IsNotFileOrDirError = OutputIsNotFileOrDirError
 
@@ -59,6 +61,7 @@ class OutputBase(object):
         cache=True,
         metric=False,
         persist=False,
+        tags=None,
     ):
         self.stage = stage
         self.repo = stage.repo
@@ -68,6 +71,7 @@ class OutputBase(object):
         self.use_cache = False if self.IS_DEPENDENCY else cache
         self.metric = False if self.IS_DEPENDENCY else metric
         self.persist = persist
+        self.tags = None if self.IS_DEPENDENCY else (tags or {})
 
         if (
             self.use_cache
@@ -197,6 +201,9 @@ class OutputBase(object):
         ret[self.PARAM_METRIC] = self.metric
         ret[self.PARAM_PERSIST] = self.persist
 
+        if self.tags:
+            ret[self.PARAM_TAGS] = self.tags
+
         return ret
 
     def verify_metric(self):
@@ -207,12 +214,20 @@ class OutputBase(object):
     def download(self, to_info, resume=False):
         self.remote.download([self.path_info], [to_info], resume=resume)
 
-    def checkout(self, force=False, progress_callback=None):
+    def checkout(self, force=False, progress_callback=None, tag=None):
         if not self.use_cache:
             return
 
+        if tag:
+            info = self.tags[tag]
+        else:
+            info = self.info
+
         getattr(self.repo.cache, self.scheme).checkout(
-            self.path_info, self.info, force=force, progress_callback=progress_callback
+            self.path_info,
+            info,
+            force=force,
+            progress_callback=progress_callback,
         )
 
     def remove(self, ignore_remove=False):

dvc/output/hdfs.py

@@ -19,6 +19,7 @@ class OutputHDFS(OutputBase):
         cache=True,
         metric=False,
         persist=False,
+        tags=None,
     ):
         super(OutputHDFS, self).__init__(
             stage,
@@ -28,6 +29,7 @@ class OutputHDFS(OutputBase):
             cache=cache,
             metric=metric,
             persist=persist,
+            tags=tags,
         )
         if remote:
             path = posixpath.join(remote.url, urlparse(path).path.lstrip("/"))

dvc/output/local.py

@@ -22,6 +22,7 @@ class OutputLOCAL(OutputBase):
         cache=True,
         metric=False,
         persist=False,
+        tags=None,
     ):
         super(OutputLOCAL, self).__init__(
             stage,
@@ -31,6 +32,7 @@ class OutputLOCAL(OutputBase):
             cache=cache,
             metric=metric,
             persist=persist,
+            tags=tags,
         )
         if remote:
             p = os.path.join(

dvc/output/s3.py

@@ -19,6 +19,7 @@ class OutputS3(OutputBase):
         cache=True,
         metric=False,
         persist=False,
+        tags=None,
     ):
         super(OutputS3, self).__init__(
             stage,
@@ -28,6 +29,7 @@ class OutputS3(OutputBase):
             cache=cache,
             metric=metric,
             persist=persist,
+            tags=tags,
         )
         bucket = remote.bucket if remote else urlparse(path).netloc
         path = urlparse(path).path.lstrip("/")

dvc/output/ssh.py

@@ -20,6 +20,7 @@ class OutputSSH(OutputBase):
         cache=True,
         metric=False,
         persist=False,
+        tags=None,
     ):
         super(OutputSSH, self).__init__(
             stage,
@@ -29,6 +30,7 @@ class OutputSSH(OutputBase):
             cache=cache,
             metric=metric,
             persist=persist,
+            tags=tags,
         )
         parsed = urlparse(path)
         host = remote.host if remote else parsed.hostname

dvc/repo/__init__.py

@@ -45,6 +45,7 @@ class Repo(object):
         from dvc.updater import Updater
         from dvc.repo.metrics import Metrics
         from dvc.scm.tree import WorkingTree
+        from dvc.repo.tag import Tag
 
         root_dir = self.find_root(root_dir)
 
@@ -70,6 +71,7 @@ class Repo(object):
         self.updater = Updater(self.dvc_dir)
 
         self.metrics = Metrics(self)
+        self.tag = Tag(self)
 
         self._ignore()
 

dvc/repo/tag/__init__.py

@@ -0,0 +1,18 @@
+class Tag(object):
+    def __init__(self, repo):
+        self.repo = repo
+
+    def add(self, *args, **kwargs):
+        from dvc.repo.tag.add import add
+
+        return add(self.repo, *args, **kwargs)
+
+    def list(self, *args, **kwargs):
+        from dvc.repo.tag.list import list
+
+        return list(self.repo, *args, **kwargs)
+
+    def remove(self, *args, **kwargs):
+        from dvc.repo.tag.remove import remove
+
+        return remove(self.repo, *args, **kwargs)

dvc/repo/tag/add.py

@@ -0,0 +1,15 @@
+import dvc.logger as logger
+
+
+def add(self, tag, target=None, with_deps=False, recursive=False):
+    stages = self.collect(target, with_deps=with_deps, recursive=recursive)
+    for stage in stages:
+        changed = False
+        for out in stage.outs:
+            if not out.info:
+                logger.warning("missing checksum info for '{}'".format(out))
+                continue
+            out.tags[tag] = out.info.copy()
+            changed = True
+        if changed:
+            stage.dump()

dvc/repo/tag/list.py

@@ -0,0 +1,11 @@
+def list(self, target=None, with_deps=False, recursive=False):
+    ret = {}
+    stages = self.collect(target, with_deps=with_deps, recursive=recursive)
+    for stage in stages:
+        outs_tags = {}
+        for out in stage.outs:
+            if out.tags:
+                outs_tags[str(out)] = out.tags.copy()
+        if outs_tags:
+            ret.update({stage.relpath: outs_tags})
+    return ret

dvc/repo/tag/remove.py

@@ -0,0 +1,15 @@
+import dvc.logger as logger
+
+
+def remove(self, tag, target=None, with_deps=False, recursive=False):
+    stages = self.collect(target, with_deps=with_deps, recursive=recursive)
+    for stage in stages:
+        changed = False
+        for out in stage.outs:
+            if tag not in out.tags.keys():
+                logger.warning("tag '{}' not found for '{}'".format(tag, out))
+                continue
+            del out.tags[tag]
+            changed = True
+        if changed:
+            stage.dump()

dvc/stage.py

@@ -2,6 +2,7 @@ from __future__ import unicode_literals
 
 from dvc.utils.compat import str, open
 
+import re
 import os
 import yaml
 import subprocess
@@ -129,6 +130,8 @@ class Stage(object):
         Optional(PARAM_LOCKED): bool,
     }
 
+    TAG_REGEX = r"^(?P<path>.*)@(?P<tag>[^\\/@:]*)$"
+
     def __init__(
         self,
         repo,
@@ -139,6 +142,7 @@ class Stage(object):
         outs=None,
         md5=None,
         locked=False,
+        tag=None,
     ):
         if deps is None:
             deps = []
@@ -153,6 +157,7 @@ class Stage(object):
         self.deps = deps
         self.md5 = md5
         self.locked = locked
+        self.tag = tag
 
     def __repr__(self):
         return "Stage: '{path}'".format(
@@ -536,8 +541,17 @@ class Stage(object):
         if not repo.tree.isfile(fname):
             raise StageFileIsNotDvcFileError(fname)
 
+    @classmethod
+    def _get_path_tag(cls, s):
+        regex = re.compile(cls.TAG_REGEX)
+        match = regex.match(s)
+        if not match:
+            return s, None
+        return match.group("path"), match.group("tag")
+
     @staticmethod
     def load(repo, fname):
+        fname, tag = Stage._get_path_tag(fname)
 
         # it raises the proper exceptions by priority:
         # 1. when the file doesn't exists
@@ -563,6 +577,7 @@ class Stage(object):
             cmd=d.get(Stage.PARAM_CMD),
             md5=d.get(Stage.PARAM_MD5),
             locked=d.get(Stage.PARAM_LOCKED, False),
+            tag=tag,
         )
 
         stage.deps = dependency.loadd_from(stage, d.get(Stage.PARAM_DEPS, []))
@@ -624,7 +639,14 @@ class Stage(object):
         # NOTE: excluding parameters that don't affect the state of the
         # pipeline. Not excluding `OutputLOCAL.PARAM_CACHE`, because if
         # it has changed, we might not have that output in our cache.
-        m = dict_md5(d, exclude=[self.PARAM_LOCKED, OutputLOCAL.PARAM_METRIC])
+        m = dict_md5(
+            d,
+            exclude=[
+                self.PARAM_LOCKED,
+                OutputLOCAL.PARAM_METRIC,
+                OutputLOCAL.PARAM_TAGS,
+            ],
+        )
         logger.debug("Computed stage '{}' md5: '{}'".format(self.relpath, m))
         return m
 
@@ -791,7 +813,9 @@ class Stage(object):
 
     def checkout(self, force=False, progress_callback=None):
         for out in self.outs:
-            out.checkout(force=force, progress_callback=progress_callback)
+            out.checkout(
+                force=force, tag=self.tag, progress_callback=progress_callback
+            )
 
     @staticmethod
     def _status(entries):