diff --git a/dvc/exceptions.py b/dvc/exceptions.py index d5faf6268..3176bf900 100644 --- a/dvc/exceptions.py +++ b/dvc/exceptions.py @@ -148,3 +148,11 @@ class CyclicGraphError(DvcException): "{stages}".format(stages=stages) ) super(CyclicGraphError, self).__init__(msg) + + +class ConfirmRemoveError(DvcException): + def __init__(self, path): + super(ConfirmRemoveError, self).__init__( + "unable to remove '{}' without a confirmation from the user. Use " + "'-f' to force.".format(path) + ) diff --git a/dvc/prompt.py b/dvc/prompt.py index 3d07e08b9..6853ce40b 100644 --- a/dvc/prompt.py +++ b/dvc/prompt.py @@ -14,7 +14,10 @@ def _ask(prompt, limited_to=None): return None while True: - answer = input(prompt + " ").lower() + try: + answer = input(prompt + " ").lower() + except EOFError: + return None if not limited_to: return answer diff --git a/dvc/remote/base.py b/dvc/remote/base.py index 6d2859c9c..ff1395c48 100644 --- a/dvc/remote/base.py +++ b/dvc/remote/base.py @@ -11,7 +11,7 @@ from multiprocessing import cpu_count import dvc.prompt as prompt import dvc.logger as logger from dvc.config import Config -from dvc.exceptions import DvcException +from dvc.exceptions import DvcException, ConfirmRemoveError STATUS_OK = 1 @@ -296,10 +296,7 @@ class RemoteBase(object): ) if not prompt.confirm(msg): - raise DvcException( - "unable to remove '{}' without a confirmation" - " from the user. Use '-f' to force.".format(str(path_info)) - ) + raise ConfirmRemoveError(str(path_info)) self.remove(path_info) diff --git a/tests/test_checkout.py b/tests/test_checkout.py index a1bae46ec..055287501 100644 --- a/tests/test_checkout.py +++ b/tests/test_checkout.py @@ -1,4 +1,5 @@ import os +import mock import yaml import shutil import filecmp @@ -11,7 +12,7 @@ from tests.basic_env import TestDvc from tests.test_repro import TestRepro from dvc.stage import Stage from dvc.remote.local import RemoteLOCAL -from dvc.exceptions import DvcException +from dvc.exceptions import DvcException, ConfirmRemoveError from mock import patch @@ -345,3 +346,21 @@ class TestCheckoutDirectory(TestRepro): self.assertEqual(ret, 0) self.assertTrue(os.path.exists(self.DATA_DIR)) + + +class TestCheckoutHook(TestDvc): + @mock.patch("sys.stdout.isatty", return_value=True) + @mock.patch("dvc.utils.compat.input", side_effect=EOFError) + def test(self, mock_input, mock_isatty): + """ Test that dvc checkout handles EOFError gracefully, which is what + it will experience when running in a git hook. + """ + stages = self.dvc.add(self.DATA_DIR) + self.assertEqual(len(stages), 1) + stage = stages[0] + self.assertNotEqual(stage, None) + + self.create(os.path.join(self.DATA_DIR, "test"), "test") + + with self.assertRaises(ConfirmRemoveError): + self.dvc.checkout() diff --git a/tests/unit/prompt.py b/tests/unit/prompt.py new file mode 100644 index 000000000..2b6395ac7 --- /dev/null +++ b/tests/unit/prompt.py @@ -0,0 +1,14 @@ +import mock +from unittest import TestCase + +from dvc.prompt import confirm + + +class TestConfirm(TestCase): + @mock.patch("sys.stdout.isatty", return_value=True) + @mock.patch("dvc.prompt.input", side_effect=EOFError) + def test_eof(self, mock_input, mock_isatty): + ret = confirm("message") + mock_isatty.assert_called() + mock_input.assert_called() + self.assertFalse(ret)