Merge pull request #1671 from mroutis/fix-1356

add: support the `--file` option

dvc/command/add.py

@@ -7,16 +7,21 @@ from dvc.command.base import CmdBase
 
 class CmdAdd(CmdBase):
     def run(self):
-        for target in self.args.targets:
-            try:
+        try:
+            if len(self.args.targets) > 1 and self.args.file:
+                raise DvcException("can't use '--file' with multiple targets")
+
+            for target in self.args.targets:
                 self.repo.add(
                     target,
                     recursive=self.args.recursive,
                     no_commit=self.args.no_commit,
+                    fname=self.args.file,
                 )
-            except DvcException:
-                logger.error("failed to add file")
-                return 1
+
+        except DvcException:
+            logger.error("failed to add file")
+            return 1
         return 0
 
 
@@ -38,6 +43,16 @@ def add_parser(subparsers, parent_parser):
         default=False,
         help="Don't put files/directories into cache.",
     )
+    add_parser.add_argument(
+        "-f",
+        "--file",
+        help="Specify name of the stage file. It should be "
+        "either 'Dvcfile' or have a '.dvc' suffix (e.g. "
+        "'prepare.dvc', 'clean.dvc', etc) in order for "
+        "dvc to be able to find it later. By default "
+        "the output basename + .dvc is used as a stage filename. "
+        "(NOTE: It can't be used when specifying multiple targets)",
+    )
     add_parser.add_argument(
         "targets", nargs="+", help="Input files/directories."
     )

dvc/exceptions.py

@@ -212,8 +212,14 @@ class NoMetricsError(DvcException):
 
 class StageFileCorruptedError(DvcException):
     def __init__(self, path):
-        msg = (
-            "Unable to read stage file: {} "
-            "\n Yaml file structure is corrupted".format(path)
+        super(StageFileCorruptedError, self).__init__(
+            "unable to read stage file: {} "
+            "YAML file structure is corrupted".format(path)
+        )
+
+
+class RecursiveAddingWhileUsingFilename(DvcException):
+    def __init__(self):
+        super(RecursiveAddingWhileUsingFilename, self).__init__(
+            "using fname with recursive is not allowed."
         )
-        super(StageFileCorruptedError, self).__init__(msg)

dvc/repo/add.py

@@ -1,11 +1,25 @@
 import os
 from dvc.stage import Stage
+from dvc.utils import walk_files
+from dvc.exceptions import RecursiveAddingWhileUsingFilename
 
 
-def add(repo, fname, recursive=False, no_commit=False):
-    fnames = _collect_filenames_to_add(repo, fname, recursive)
+def add(repo, target, recursive=False, no_commit=False, fname=None):
+    if recursive and fname:
+        raise RecursiveAddingWhileUsingFilename()
 
-    stages = _create_stagefiles(repo, fnames, no_commit)
+    targets = [target]
+
+    if os.path.isdir(target) and recursive:
+        targets = [
+            file
+            for file in walk_files(target)
+            if not Stage.is_stage_file(file)
+            if os.path.basename(file) != repo.scm.ignore_file
+            if not repo.scm.is_tracked(file)
+        ]
+
+    stages = _create_stages(repo, targets, fname, no_commit)
 
     repo.check_dag(repo.stages() + stages)
 
@@ -17,49 +31,21 @@ def add(repo, fname, recursive=False, no_commit=False):
     return stages
 
 
-def _collect_filenames_to_add(repo, fname, recursive):
-    if recursive and os.path.isdir(fname):
-        fnames = _collect_valid_filenames_from_directory(repo, fname)
-    else:
-        fnames = [fname]
-    return fnames
-
-
-def _collect_valid_filenames_from_directory(repo, fname):
-    fnames = []
-    for file_path in _file_paths(fname):
-        if _is_valid_file_to_add(file_path, repo):
-            fnames.append(file_path)
-    return fnames
-
-
-def _file_paths(directory):
-    for root, _, files in os.walk(directory):
-        for f in files:
-            yield os.path.join(root, f)
-
-
-def _is_valid_file_to_add(file_path, repo):
-    if Stage.is_stage_file(file_path):
-        return False
-    if os.path.basename(file_path) == repo.scm.ignore_file:
-        return False
-    if repo.scm.is_tracked(file_path):
-        return False
-    return True
-
-
-def _create_stagefiles(repo, fnames, no_commit):
+def _create_stages(repo, targets, fname, no_commit):
     stages = []
-    with repo.state:
-        for f in fnames:
-            stage = Stage.create(repo=repo, outs=[f], add=True)
 
-            if stage is None:
+    with repo.state:
+        for out in targets:
+            stage = Stage.create(repo=repo, outs=[out], add=True, fname=fname)
+
+            if not stage:
                 continue
 
             stage.save()
+
             if not no_commit:
                 stage.commit()
+
             stages.append(stage)
+
     return stages

dvc/utils/__init__.py

@@ -239,3 +239,9 @@ def load_stage_file(path):
             return yaml.safe_load(fobj) or {}
         except ScannerError:
             raise StageFileCorruptedError(path)
+
+
+def walk_files(directory):
+    for root, _, files in os.walk(directory):
+        for f in files:
+            yield os.path.join(root, f)

tests/test_add.py

@@ -16,7 +16,7 @@ from mock import patch
 from dvc.main import main
 from dvc.utils import file_md5, load_stage_file
 from dvc.stage import Stage
-from dvc.exceptions import DvcException
+from dvc.exceptions import DvcException, RecursiveAddingWhileUsingFilename
 from dvc.output.base import OutputAlreadyTrackedError
 from dvc.repo import Repo as DvcRepo
 
@@ -425,7 +425,35 @@ class TestShouldThrowProperExceptionOnCorruptedStageFile(TestDvc):
             self.assertEqual(1, ret)
 
             self.assertIn(
-                "Unable to read stage file: {} \n "
-                "Yaml file structure is corrupted".format(foo_stage),
+                "unable to read stage file: {} "
+                "YAML file structure is corrupted".format(foo_stage),
                 logger.handlers[1].stream.getvalue(),
             )
+
+
+class TestAddFilename(TestDvc):
+    def test(self):
+        ret = main(["add", self.FOO, self.BAR, "-f", "error.dvc"])
+        self.assertNotEqual(0, ret)
+
+        ret = main(["add", "-R", self.DATA_DIR, "-f", "error.dvc"])
+        self.assertNotEqual(0, ret)
+
+        with self.assertRaises(RecursiveAddingWhileUsingFilename):
+            self.dvc.add(self.DATA_DIR, recursive=True, fname="error.dvc")
+
+        ret = main(["add", self.DATA_DIR, "-f", "data_directory.dvc"])
+        self.assertEqual(0, ret)
+        self.assertTrue(os.path.exists("data_directory.dvc"))
+
+        ret = main(["add", self.FOO, "-f", "bar.dvc"])
+        self.assertEqual(0, ret)
+        self.assertTrue(os.path.exists("bar.dvc"))
+        self.assertFalse(os.path.exists("foo.dvc"))
+
+        os.remove("bar.dvc")
+
+        ret = main(["add", self.FOO, "--file", "bar.dvc"])
+        self.assertEqual(0, ret)
+        self.assertTrue(os.path.exists("bar.dvc"))
+        self.assertFalse(os.path.exists("foo.dvc"))