From 4271ddf11cc991dd81f3fdcd9b52c317ee38334c Mon Sep 17 00:00:00 2001 From: Ruslan Kuprieiev Date: Fri, 15 Jun 2018 18:37:50 +0300 Subject: [PATCH] dvc introduce `dvc import` Fixes #705 Signed-off-by: Ruslan Kuprieiev --- dvc/cli.py | 20 +++++++++++++++++--- dvc/command/imp.py | 12 ++++++++++++ dvc/dependency/base.py | 3 +++ dvc/project.py | 15 +++++++++++++++ dvc/remote/base.py | 6 ++++++ dvc/remote/gs.py | 16 +++++++++++++++- dvc/remote/hdfs.py | 17 +++++++++++++++++ dvc/remote/local.py | 19 +++++++++++++++++++ dvc/remote/s3.py | 13 +++++++++++++ dvc/stage.py | 33 +++++++++++++++++++++++---------- tests/test_repro.py | 12 ++++++++++-- 11 files changed, 150 insertions(+), 16 deletions(-) create mode 100644 dvc/command/imp.py diff --git a/dvc/cli.py b/dvc/cli.py index 97d0c0f3f..c25867eba 100644 --- a/dvc/cli.py +++ b/dvc/cli.py @@ -14,6 +14,7 @@ from dvc.command.data_sync import CmdDataPush, CmdDataPull, CmdDataFetch from dvc.command.status import CmdDataStatus from dvc.command.gc import CmdGC from dvc.command.add import CmdAdd +from dvc.command.imp import CmdImport from dvc.command.config import CmdConfig from dvc.command.checkout import CmdCheckout from dvc.command.remote import CmdRemoteAdd, CmdRemoteRemove, CmdRemoteModify, CmdRemoteList @@ -83,15 +84,28 @@ def parse_args(argv=None): init_parser.set_defaults(func=CmdInit) # Add - import_parser = subparsers.add_parser( + add_parser = subparsers.add_parser( 'add', parents=[parent_parser], help='Add files/directories to dvc') - import_parser.add_argument( + add_parser.add_argument( 'targets', nargs='+', help='Input files/directories') - import_parser.set_defaults(func=CmdAdd) + add_parser.set_defaults(func=CmdAdd) + + # Import + import_parser = subparsers.add_parser( + 'import', + parents=[parent_parser], + help='Import files from URL') + import_parser.add_argument( + 'url', + help='URL') + import_parser.add_argument( + 'out', + help='Output') + import_parser.set_defaults(func=CmdImport) # Checkout checkout_parser = subparsers.add_parser( diff --git a/dvc/command/imp.py b/dvc/command/imp.py new file mode 100644 index 000000000..f5119e7be --- /dev/null +++ b/dvc/command/imp.py @@ -0,0 +1,12 @@ +from dvc.exceptions import DvcException +from dvc.command.common.base import CmdBase + + +class CmdImport(CmdBase): + def run(self): + try: + self.project.imp(self.args.url, self.args.out) + except DvcException as ex: + self.project.logger.error('Failed to import {}', ex) + return 1 + return 0 diff --git a/dvc/dependency/base.py b/dvc/dependency/base.py index 525e40209..ba993b21d 100644 --- a/dvc/dependency/base.py +++ b/dvc/dependency/base.py @@ -53,3 +53,6 @@ class DependencyBase(object): def dumpd(self): return {self.PARAM_PATH: self.path} + + def download(self, path): + self.remote.download(self.path_info, path) diff --git a/dvc/project.py b/dvc/project.py index f420d792d..022009751 100644 --- a/dvc/project.py +++ b/dvc/project.py @@ -159,6 +159,21 @@ class Project(object): stage.dump() return stage + def imp(self, url, out): + stage_fname = out + Stage.STAGE_FILE_SUFFIX + cwd = os.path.dirname(os.path.abspath(out)) + stage = Stage.loads(project=self, + cmd=None, + deps=[url], + outs=[out], + fname=stage_fname, + cwd=cwd) + + stage.run() + stage.dump() + return stage + + def _reproduce_stage(self, stages, node, force): stage = stages[node].reproduce(force=force) if not stage: diff --git a/dvc/remote/base.py b/dvc/remote/base.py index 721598c1e..9d05ad800 100644 --- a/dvc/remote/base.py +++ b/dvc/remote/base.py @@ -26,3 +26,9 @@ class RemoteBase(object): def checkout(self, path_info, checksum_info): pass + + def download(self, path_info, path): + pass + + def upload(self, path, path_info): + pass diff --git a/dvc/remote/gs.py b/dvc/remote/gs.py index 9cad81fed..44e5c7715 100644 --- a/dvc/remote/gs.py +++ b/dvc/remote/gs.py @@ -79,8 +79,22 @@ class RemoteGS(RemoteBase): if path_info['scheme'] != 'gs': raise NotImplementedError - blob = self.gs.bucket(path_info['key']).get_blob(path_info['key']) + blob = self.gs.bucket(path_info['bucket']).get_blob(path_info['key']) if not blob: return blob.delete() + + def upload(self, path, path_info): + if path_info['scheme'] != 'gs': + raise NotImplementedError + + blob = self.gs.bucket(path_info['bucket']).blob(path_info['key']) + blob.upload_from_filename(path) + + def download(self, path_info, path): + if path_info['scheme'] != 'gs': + raise NotImplementedError + + blob = self.gs.bucket(path_info['bucket']).get_blob(path_info['key']) + blob.download_to_filename(path) diff --git a/dvc/remote/hdfs.py b/dvc/remote/hdfs.py index 3f31b41cf..445297747 100644 --- a/dvc/remote/hdfs.py +++ b/dvc/remote/hdfs.py @@ -101,3 +101,20 @@ class RemoteHDFS(RemoteBase): assert path_info.get('url') self.rm(path_info['url']) + + def upload(self, path, path_info): + if path_info['scheme'] != 'hdfs': + raise NotImplementedError + + self.hadoop_fs('mkdir -p {}'.format(posixpath.dirname(path_info['url']))) + self.hadoop_fs('copyFromLocal {} {}'.format(path, path_info['url'])) + + def download(self, path_info, path): + if path_info['scheme'] != 'hdfs': + raise NotImplementedError + + dname = os.path.dirname(path) + if not os.path.exists(dname): + os.makedirs(dname) + + self.hadoop_fs('copyToLocal {} {}'.format(path_info['url'], path)) diff --git a/dvc/remote/local.py b/dvc/remote/local.py index 38901c215..24c00c4ba 100644 --- a/dvc/remote/local.py +++ b/dvc/remote/local.py @@ -243,3 +243,22 @@ class RemoteLOCAL(RemoteBase): raise NotImplementedError remove(path_info['path']) + + def _copy(self, from_path, to_path): + dname = os.path.dirname(to_path) + if not os.path.exists(dname): + os.makedirs(dname) + + shutil.copyfile(from_path, to_path) + + def upload(self, path, path_info): + if path_info['scheme'] != 'local': + raise NotImplementedError + + self._copy(path, path_info['path']) + + def download(self, path_info, path): + if path_info['scheme'] != 'local': + raise NotImplementedError + + self._copy(path_info['path'], path) diff --git a/dvc/remote/s3.py b/dvc/remote/s3.py index 83cce0b15..da5cb951f 100644 --- a/dvc/remote/s3.py +++ b/dvc/remote/s3.py @@ -84,3 +84,16 @@ class RemoteS3(RemoteBase): obj.delete() except Exception: pass + + def upload(self, path, path_info): + if path_info['scheme'] != 's3': + raise NotImplementedError + + self.s3.Object(path_info['bucket'], path_info['key']).upload_file(path) + + + def download(self, path_info, path): + if path_info['scheme'] != 's3': + raise NotImplementedError + + self.s3.Object(path_info['bucket'], path_info['key']).download_file(path) diff --git a/dvc/stage.py b/dvc/stage.py index 231677923..fb2d9c9f5 100644 --- a/dvc/stage.py +++ b/dvc/stage.py @@ -95,6 +95,13 @@ class Stage(object): def is_callback(self): return not self.is_data_source and len(self.deps) == 0 + @property + def is_import(self): + return not self.cmd and \ + len(self.deps) == 1 and \ + len(self.outs) == 1 and \ + self.outs[0].path_info['scheme'] == 'local' + def changed(self): ret = False @@ -127,12 +134,16 @@ class Stage(object): if not self.changed() and not force: return None - if self.cmd: + if self.cmd or self.is_import: # Removing outputs only if we actually have command to reproduce self.remove_outs(ignore_remove=False) + self.project.logger.info(u'Reproducing \'{}\''.format(self.relpath)) + self.run() + self.project.logger.debug(u'\'{}\' was reproduced'.format(self.relpath)) + return self @staticmethod @@ -223,8 +234,16 @@ class Stage(object): out.save() def run(self): - if not self.is_data_source: - self.project.logger.info(u'Reproducing \'{}\':\n\t{}'.format(self.relpath, self.cmd)) + if self.is_import: + msg = u'Importing \'{}\' -> \'{}\'' + self.project.logger.info(msg.format(self.deps[0].path, self.outs[0].rel_path)) + + self.deps[0].download(self.outs[0].path) + elif self.is_data_source: + self.project.logger.info(u'Verifying data sources in \'{}\''.format(self.relpath)) + self.check_missing_outputs() + else: + self.project.logger.info(u'Running command:\n\t{}'.format(self.cmd)) p = subprocess.Popen(self.cmd, cwd=self.cwd, @@ -234,14 +253,8 @@ class Stage(object): p.communicate() if p.returncode != 0: raise StageCmdFailedError(self) - - self.save() - self.project.logger.debug(u'\'{}\' was reproduced'.format(self.relpath)) - else: - self.project.logger.info(u'Verifying data sources in \'{}\''.format(self.relpath)) - self.check_missing_outputs() - self.save() + self.save() def check_missing_outputs(self): missing_outs = [out.rel_path for out in self.outs if not os.path.exists(out.rel_path)] diff --git a/tests/test_repro.py b/tests/test_repro.py index ff55fc1c0..64129fe22 100644 --- a/tests/test_repro.py +++ b/tests/test_repro.py @@ -271,14 +271,22 @@ class TestReproExternalBase(TestDvc): self.write(self.bucket, foo_key, 'foo') - stage = self.dvc.run(outs=[out_bar_path], + import_stage = self.dvc.imp(out_foo_path, 'import') + self.assertTrue(os.path.exists('import')) + self.assertTrue(filecmp.cmp('import', self.FOO)) + + cmd_stage = self.dvc.run(outs=[out_bar_path], deps=[out_foo_path], cmd='{} {} {}'.format(self.cmd, foo_path, bar_path)) self.write(self.bucket, foo_key, 'bar') - stages = self.dvc.reproduce(stage.path) + stages = self.dvc.reproduce(import_stage.path) + self.assertEqual(len(stages), 1) + self.assertTrue(os.path.exists('import')) + self.assertTrue(filecmp.cmp('import', self.BAR)) + stages = self.dvc.reproduce(cmd_stage.path) self.assertEqual(len(stages), 1)