From 78e5b166fee89adf80be75926be2f79a928f5dff Mon Sep 17 00:00:00 2001
From: Ruslan Kuprieiev <kupruser@gmail.com>
Date: Wed, 13 Jun 2018 13:53:06 +0300
Subject: [PATCH] dvc: replace StageNotFoundError with NotDvcFileError

Fixes #760

Signed-off-by: Ruslan Kuprieiev <kupruser@gmail.com>
---
 dvc/command/repro.py | 13 +------------
 dvc/project.py       | 16 ++++++++++------
 tests/test_add.py    |  1 -
 tests/test_gc.py     |  1 -
 tests/test_remove.py |  4 ++--
 5 files changed, 13 insertions(+), 22 deletions(-)

diff --git a/dvc/command/repro.py b/dvc/command/repro.py
index 681df032e..b5e663514 100644
--- a/dvc/command/repro.py
+++ b/dvc/command/repro.py
@@ -1,7 +1,7 @@
 import os
 
 from dvc.command.common.base import CmdBase
-from dvc.project import ReproductionError, StageNotFoundError
+from dvc.project import ReproductionError
 from dvc.stage import Stage
 
 
@@ -23,17 +23,6 @@ class CmdRepro(CmdBase):
                 self.project.logger.error(msg, ex)
                 ret = 1
                 break
-            except StageNotFoundError as ex:
-                if os.path.exists(target):
-                    msg = '\'{}\' is not a dvc file.'.format(target)
-                    dvcfile = target + Stage.STAGE_FILE_SUFFIX
-                    if os.path.exists(target) and os.path.isfile(dvcfile):
-                        msg += ' Maybe you meant \'{}\'?'.format(dvcfile)
-                else:
-                    msg = '\'{}\' does not exist.'.format(target)
-                self.project.logger.error(msg)
-                ret = 1
-                break
             except Exception as ex:
                 msg = 'Failed to reproduce \'{}\' - unexpected error'.format(target)
                 self.project.logger.error(msg, ex)
diff --git a/dvc/project.py b/dvc/project.py
index 05a6c3d36..f420d792d 100644
--- a/dvc/project.py
+++ b/dvc/project.py
@@ -25,14 +25,18 @@ class InitError(DvcException):
         super(InitError, self).__init__(msg)
 
 
-class StageNotFoundError(DvcException):
+class NotDvcFileError(DvcException):
     def __init__(self, path):
-        msg = 'Stage file {} does not exist'.format(path)
-        super(StageNotFoundError, self).__init__(msg)
+        msg = '\'{}\' is not a DVC file'.format(path)
+        p = path + Stage.STAGE_FILE_SUFFIX
+        if Stage.is_stage_file(p):
+            msg += '. Maybe you meant \'{}\'?'.format(p)
+        super(NotDvcFileError, self).__init__(msg)
 
 
 class ReproductionError(DvcException):
     def __init__(self, dvc_file_name, ex):
+        self.path = dvc_file_name
         msg = 'Failed to reproduce \'{}\''.format(dvc_file_name)
         super(ReproductionError, self).__init__(msg, cause=ex)
 
@@ -125,7 +129,7 @@ class Project(object):
 
     def remove(self, target, outs_only=False):
         if not Stage.is_stage_file(target):
-            raise StageNotFoundError(target)
+            raise NotDvcFileError(target)
 
         stage = Stage.load(self, target)
         if outs_only:
@@ -166,7 +170,7 @@ class Project(object):
         stages = nx.get_node_attributes(self.graph(), 'stage')
         node = os.path.relpath(os.path.abspath(target), self.root_dir)
         if node not in stages:
-            raise StageNotFoundError(target)
+            raise NotDvcFileError(target)
 
         if recursive:
             return self._reproduce_stages(stages, node, force)
@@ -185,7 +189,7 @@ class Project(object):
     def checkout(self, target=None):
         if target:
             if not Stage.is_stage_file(target):
-                raise StageNotFoundError(target)
+                raise NotDvcFileError(target)
             stages = [Stage.load(self, target)]
         else:
             self.link_state.remove_all()
diff --git a/tests/test_add.py b/tests/test_add.py
index b1f0736d3..0acf3194d 100644
--- a/tests/test_add.py
+++ b/tests/test_add.py
@@ -8,7 +8,6 @@ from dvc.utils import file_md5
 from dvc.stage import Stage
 from dvc.output.base import OutputAlreadyTrackedError
 from dvc.output.base import OutputDoesNotExistError, OutputIsNotFileOrDirError
-from dvc.project import StageNotFoundError
 from dvc.command.add import CmdAdd
 
 from tests.basic_env import TestDvc
diff --git a/tests/test_gc.py b/tests/test_gc.py
index 53ac95f9e..6af35b1ba 100644
--- a/tests/test_gc.py
+++ b/tests/test_gc.py
@@ -5,7 +5,6 @@ import filecmp
 from dvc.main import main
 from dvc.utils import file_md5
 from dvc.stage import Stage
-from dvc.project import StageNotFoundError
 from dvc.command.gc import CmdGC
 
 from tests.basic_env import TestDvc
diff --git a/tests/test_remove.py b/tests/test_remove.py
index 7ae68f6ee..4e2ad1ef5 100644
--- a/tests/test_remove.py
+++ b/tests/test_remove.py
@@ -2,7 +2,7 @@ import os
 
 from dvc.main import main
 from dvc.stage import Stage
-from dvc.project import StageNotFoundError
+from dvc.project import NotDvcFileError
 from dvc.command.remove import CmdRemove
 
 from tests.basic_env import TestDvc
@@ -26,7 +26,7 @@ class TestRemove(TestDvc):
 
 class TestRemoveNonExistentFile(TestDvc):
     def test(self):
-        with self.assertRaises(StageNotFoundError) as cx:
+        with self.assertRaises(NotDvcFileError) as cx:
             self.dvc.remove('non_existent_dvc_file')