diff --git a/dvc/command/add.py b/dvc/command/add.py index eb0ac3af4..68c81db52 100644 --- a/dvc/command/add.py +++ b/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." ) diff --git a/dvc/exceptions.py b/dvc/exceptions.py index b225d61f8..c27bf3837 100644 --- a/dvc/exceptions.py +++ b/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) diff --git a/dvc/repo/add.py b/dvc/repo/add.py index 87916f94c..b85883f6d 100644 --- a/dvc/repo/add.py +++ b/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 diff --git a/dvc/utils/__init__.py b/dvc/utils/__init__.py index b7b02d4ad..d45b13cbd 100644 --- a/dvc/utils/__init__.py +++ b/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) diff --git a/tests/test_add.py b/tests/test_add.py index af2ee388d..447b11d16 100644 --- a/tests/test_add.py +++ b/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"))