From b07c386d0202d28ebe84d83c372f30534885393f Mon Sep 17 00:00:00 2001
From: vyloy <vyloy.w@gmail.com>
Date: Mon, 15 Apr 2019 14:38:49 +0800
Subject: [PATCH] import: support the '--file' option

---
 dvc/command/imp.py        | 14 +++++++++++++-
 dvc/repo/imp.py           |  6 ++++--
 tests/func/test_import.py | 20 ++++++++++++++++++++
 3 files changed, 37 insertions(+), 3 deletions(-)

diff --git a/dvc/command/imp.py b/dvc/command/imp.py
index 726bb193c..3cd992bac 100644
--- a/dvc/command/imp.py
+++ b/dvc/command/imp.py
@@ -19,7 +19,9 @@ class CmdImport(CmdBase):
 
             out = self.args.out or default_out
 
-            self.repo.imp(self.args.url, out, self.args.resume)
+            self.repo.imp(
+                self.args.url, out, self.args.resume, fname=self.args.file
+            )
         except DvcException:
             logger.exception(
                 "failed to import {}. You could also try downloading "
@@ -62,4 +64,14 @@ def add_parser(subparsers, parent_parser):
     import_parser.add_argument(
         "out", nargs="?", help="Destination path to put files to."
     )
+    import_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 first output basename + .dvc is used as "
+        "a stage filename.",
+    )
     import_parser.set_defaults(func=CmdImport)
diff --git a/dvc/repo/imp.py b/dvc/repo/imp.py
index e317eb008..acbbc4afd 100644
--- a/dvc/repo/imp.py
+++ b/dvc/repo/imp.py
@@ -2,11 +2,13 @@ from dvc.repo.scm_context import scm_context
 
 
 @scm_context
-def imp(self, url, out, resume=False):
+def imp(self, url, out, resume=False, fname=None):
     from dvc.stage import Stage
 
     with self.state:
-        stage = Stage.create(repo=self, cmd=None, deps=[url], outs=[out])
+        stage = Stage.create(
+            repo=self, cmd=None, deps=[url], outs=[out], fname=fname
+        )
 
     if stage is None:
         return None
diff --git a/tests/func/test_import.py b/tests/func/test_import.py
index 8d770cacd..5d47bcdc2 100644
--- a/tests/func/test_import.py
+++ b/tests/func/test_import.py
@@ -156,3 +156,23 @@ class TestShouldRemoveOutsBeforeImport(TestDvc):
             self.assertEqual(0, ret)
 
         self.assertEqual(1, remove_outs_call_counter.mock.call_count)
+
+
+class TestImportFilename(TestDvc):
+    def setUp(self):
+        super(TestImportFilename, self).setUp()
+        tmp_dir = self.mkdtemp()
+        self.external_source = os.path.join(tmp_dir, "file")
+        with open(self.external_source, "w") as fobj:
+            fobj.write("content")
+
+    def test(self):
+        ret = main(["import", "-f", "bar.dvc", self.external_source])
+        self.assertEqual(0, ret)
+        self.assertTrue(os.path.exists("bar.dvc"))
+
+        os.remove("bar.dvc")
+
+        ret = main(["import", "--file", "bar.dvc", self.external_source])
+        self.assertEqual(0, ret)
+        self.assertTrue(os.path.exists("bar.dvc"))