From fa89ec8789b0e3e68a77dbf83d6736502fc6963b Mon Sep 17 00:00:00 2001
From: Ruslan Kuprieiev <ruslan@iterative.ai>
Date: Fri, 24 Aug 2018 18:05:39 +0300
Subject: [PATCH] dvc: check stage filename

Fixes #1054

Signed-off-by: Ruslan Kuprieiev <ruslan@iterative.ai>
---
 dvc/cli.py        |  7 ++++++-
 dvc/stage.py      | 28 ++++++++++++++++++++++++----
 tests/test_run.py | 13 ++++++++++++-
 3 files changed, 42 insertions(+), 6 deletions(-)

diff --git a/dvc/cli.py b/dvc/cli.py
index b8b04be72..bec8c7abe 100644
--- a/dvc/cli.py
+++ b/dvc/cli.py
@@ -217,7 +217,12 @@ def parse_args(argv=None):
     run_parser.add_argument(
                         '-f',
                         '--file',
-                        help='Specify name of the state 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.")
     run_parser.add_argument(
                         '-c',
                         '--cwd',
diff --git a/dvc/stage.py b/dvc/stage.py
index eef50c6f8..25d6ee045 100644
--- a/dvc/stage.py
+++ b/dvc/stage.py
@@ -36,6 +36,13 @@ class StageFileIsNotDvcFileError(DvcException):
         super(StageFileIsNotDvcFileError, self).__init__(msg)
 
 
+class StageFileBadNameError(DvcException):
+    def __init__(self, fname):
+        msg = "Bad stage filename '{}'. Stage files should be named " \
+              "'Dvcfile' or have a '.dvc' suffix(e.g. '{}.dvc')."
+        super(StageFileBadNameError, self).__init__(msg.format(fname, fname))
+
+
 class MissingDataSource(DvcException):
     def __init__(self, missing_files):
         assert len(missing_files) > 0
@@ -93,16 +100,20 @@ class Stage(object):
         return self.cmd is None
 
     @staticmethod
-    def is_stage_file(path):
-        if not os.path.isfile(path):
-            return False
-
+    def is_stage_filename(path):
         if not path.endswith(Stage.STAGE_FILE_SUFFIX) \
            and os.path.basename(path) != Stage.STAGE_FILE:
             return False
 
         return True
 
+    @staticmethod
+    def is_stage_file(path):
+        if not os.path.isfile(path):
+            return False
+
+        return Stage.is_stage_filename(path)
+
     def changed_md5(self):
         md5 = self.dumpd().get(self.PARAM_MD5, None)
 
@@ -276,6 +287,11 @@ class Stage(object):
 
         return stage
 
+    @staticmethod
+    def _check_dvc_filename(fname):
+        if not Stage.is_stage_filename(fname):
+            raise StageFileBadNameError(os.path.relpath(fname))
+
     @staticmethod
     def _check_dvc_file(fname):
         sname = fname + Stage.STAGE_FILE_SUFFIX
@@ -288,6 +304,8 @@ class Stage(object):
             Stage._check_dvc_file(fname)
             raise StageFileDoesNotExistError(fname)
 
+        Stage._check_dvc_filename(fname)
+
         if not Stage.is_stage_file(fname):
             Stage._check_dvc_file(fname)
             raise StageFileIsNotDvcFileError(fname)
@@ -320,6 +338,8 @@ class Stage(object):
         if not fname:
             fname = self.path
 
+        self._check_dvc_filename(fname)
+
         msg = "Saving information to '{}'.".format(os.path.relpath(fname))
         Logger.info(msg)
 
diff --git a/tests/test_run.py b/tests/test_run.py
index bee5e996c..32052fc9e 100644
--- a/tests/test_run.py
+++ b/tests/test_run.py
@@ -3,7 +3,7 @@ import filecmp
 
 from dvc.main import main
 from dvc.utils import file_md5
-from dvc.stage import Stage
+from dvc.stage import Stage, StageFileBadNameError
 from dvc.command.run import CmdRun
 from dvc.exceptions import OutputDuplicationError
 
@@ -55,6 +55,17 @@ class TestRunEmpty(TestDvc):
                      cwd=os.curdir)
 
 
+class TestRunBadStageFilename(TestDvc):
+    def test(self):
+        with self.assertRaises(StageFileBadNameError):
+            self.dvc.run(cmd='',
+                         deps=[],
+                         outs=[],
+                         outs_no_cache=[],
+                         fname='empty',
+                         cwd=os.curdir)
+
+
 class TestRunNoExec(TestDvc):
     def test(self):
         self.dvc.run(cmd='python {} {} {}'.format(self.CODE, self.FOO, 'out'),