Merge pull request #1329 from mroutis/fix-1328

checkout: Prompt when deleting files from working dir

dvc/cli.py

@@ -192,6 +192,13 @@ def parse_args(argv=None):
                         default=False,
                         help='Checkout all dependencies of the '
                              'specified target.')
+    checkout_parser.add_argument(
+                        '-f',
+                        '--force',
+                        action='store_true',
+                        default=False,
+                        help='Do not prompt when removing '
+                             'working directory files.')
     checkout_parser.add_argument(
                         'targets',
                         nargs='*',
@@ -318,6 +325,11 @@ def parse_args(argv=None):
                         default=False,
                         help='Fetch cache for all dependencies of the '
                              'specified target.')
+    pull_parser.add_argument(
+                        '-f',
+                        '--force',
+                        help='Do not prompt when removing '
+                             'working directory files.')
     pull_parser.set_defaults(func=CmdDataPull)
 
     # Push

dvc/command/checkout.py

@@ -4,9 +4,10 @@ from dvc.command.base import CmdBase
 class CmdCheckout(CmdBase):
     def run(self):
         if not self.args.targets:
-            self.project.checkout()
+            self.project.checkout(force=self.args.force)
         else:
             for target in self.args.targets:
                 self.project.checkout(target=target,
-                                      with_deps=self.args.with_deps)
+                                      with_deps=self.args.with_deps,
+                                      force=self.args.force)
         return 0

dvc/command/data_sync.py

@@ -25,7 +25,8 @@ class CmdDataPull(CmdDataBase):
                               show_checksums=self.args.show_checksums,
                               all_branches=self.args.all_branches,
                               all_tags=self.args.all_tags,
-                              with_deps=self.args.with_deps)
+                              with_deps=self.args.with_deps,
+                              force=self.args.force)
         except Exception as exc:
             self.project.logger.error('Failed to pull data from the cloud',
                                       exc)

dvc/output/gs.py

@@ -35,7 +35,7 @@ class OutputGS(DependencyGS):
 
         return False
 
-    def checkout(self):
+    def checkout(self, force=False):
         if not self.use_cache:
             return
 

dvc/output/hdfs.py

@@ -35,7 +35,7 @@ class OutputHDFS(DependencyHDFS):
 
         return False
 
-    def checkout(self):
+    def checkout(self, force=False):
         if not self.use_cache:
             return
 

dvc/output/local.py

@@ -60,11 +60,13 @@ class OutputLOCAL(DependencyLOCAL):
 
         return self.project.cache.local.changed(self.path_info, self.info)
 
-    def checkout(self):
+    def checkout(self, force=False):
         if not self.use_cache:
             return
 
-        self.project.cache.local.checkout(self.path_info, self.info)
+        self.project.cache.local.checkout(self.path_info,
+                                          self.info,
+                                          force=force)
 
     def _verify_metric(self):
         if not self.metric:

dvc/output/s3.py

@@ -35,7 +35,7 @@ class OutputS3(DependencyS3):
 
         return False
 
-    def checkout(self):
+    def checkout(self, force=False):
         if not self.use_cache:
             return
 

dvc/output/ssh.py

@@ -35,7 +35,7 @@ class OutputSSH(DependencySSH):
 
         return False
 
-    def checkout(self):
+    def checkout(self, force=False):
         if not self.use_cache:
             return
 

dvc/project.py

@@ -466,7 +466,7 @@ class Project(object):
                 used.append(out.path)
         self.state.remove_unused_links(used)
 
-    def checkout(self, target=None, with_deps=False):
+    def checkout(self, target=None, with_deps=False, force=False):
         all_stages = self.active_stages()
         stages = all_stages
 
@@ -482,7 +482,7 @@ class Project(object):
                           'not going to be checked out.'
                     self.logger.warn(msg.format(stage.relpath))
 
-                stage.checkout()
+                stage.checkout(force=force)
 
     def _get_pipeline(self, node):
         pipelines = list(filter(lambda g: node in g.nodes(),
@@ -762,7 +762,8 @@ class Project(object):
              all_branches=False,
              show_checksums=False,
              with_deps=False,
-             all_tags=False):
+             all_tags=False,
+             force=False):
         self.fetch(target,
                    jobs,
                    remote=remote,
@@ -770,7 +771,7 @@ class Project(object):
                    all_tags=all_tags,
                    show_checksums=show_checksums,
                    with_deps=with_deps)
-        self.checkout(target=target, with_deps=with_deps)
+        self.checkout(target=target, with_deps=with_deps, force=force)
 
     def _local_status(self, target=None):
         status = {}
@@ -957,7 +958,7 @@ class Project(object):
 
             for fname, t, x in entries:
                 if stage:
-                    self.checkout(stage)
+                    self.checkout(stage, force=True)
 
                 rel = os.path.relpath(fname)
                 metric = self._read_metric(fname,

dvc/remote/base.py

@@ -92,7 +92,7 @@ class RemoteBase(object):
     def save(self, path_info):
         raise NotImplementedError
 
-    def checkout(self, path_info, checksum_info):
+    def checkout(self, path_info, checksum_info, force=False):
         raise NotImplementedError
 
     def download(self, from_infos, to_infos, no_progress_bar=False, name=None):

dvc/remote/local.py

@@ -268,7 +268,7 @@ class RemoteLOCAL(RemoteBase):
     def is_dir_cache(cls, cache):
         return cache.endswith(cls.MD5_DIR_SUFFIX)
 
-    def checkout(self, path_info, checksum_info):
+    def checkout(self, path_info, checksum_info, force=False):
         path = path_info['path']
         md5 = checksum_info.get(self.PARAM_MD5)
         cache = self.get(md5)
@@ -295,9 +295,7 @@ class RemoteLOCAL(RemoteBase):
         if not self.is_dir_cache(cache):
 
             if os.path.exists(path):
-                msg = u'File \'{}\' exists. Removing before checkout.'
-                Logger.warn(msg.format(os.path.relpath(path)))
-                remove(path)
+                remove(path) if force else self._safe_remove(path)
 
             self.link(cache, path)
             self.state.update_link(path)
@@ -331,20 +329,20 @@ class RemoteLOCAL(RemoteBase):
             }
 
             if self.changed(entry_info, entry_checksum_info):
-                remove(p)
+                remove(p) if force else self._safe_remove(p)
                 self.link(c, p)
 
             if bar:
                 progress.update_target(dir_relpath, processed, dir_size)
 
-        self._discard_working_directory_changes(path, dir_info)
+        self._discard_working_directory_changes(path, dir_info, force=force)
 
         self.state.update_link(path)
 
         if bar:
             progress.finish_target(dir_relpath)
 
-    def _discard_working_directory_changes(self, path, dir_info):
+    def _discard_working_directory_changes(self, path, dir_info, force=False):
         working_dir_files = set(
             os.path.join(root, file)
             for root, _, files in os.walk(path)
@@ -359,7 +357,21 @@ class RemoteLOCAL(RemoteBase):
         delta = working_dir_files - cached_files
 
         for file in delta:
-            remove(file)
+            remove(file) if force else self._safe_remove(file)
+
+    def _safe_remove(self, file):
+        msg = (
+            'File "{}" is going to be removed. '
+            'Are you sure you want to proceed?'
+            .format(file)
+        )
+
+        confirmed = self.project.prompt.prompt(msg, False)
+
+        if not confirmed:
+            raise DvcException('Unable to remove {} without a confirmation'
+                               " from the user. Use '-f' to force."
+                               .format(file))
 
     def _move(self, inp, outp):
         # moving in two stages to make last the move atomic in

dvc/stage.py

@@ -478,9 +478,9 @@ class Stage(object):
         if paths:
             raise MissingDataSource(paths)
 
-    def checkout(self):
+    def checkout(self, force=False):
         for out in self.outs:
-            out.checkout()
+            out.checkout(force=force)
 
     def _status(self, entries, name):
         ret = {}

tests/test_checkout.py

@@ -13,7 +13,9 @@ 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 mock import patch
 
 class TestCheckout(TestRepro):
     def setUp(self):
@@ -34,7 +36,7 @@ class TestCheckout(TestRepro):
         shutil.rmtree(self.DATA_DIR)
 
     def test(self):
-        self.dvc.checkout()
+        self.dvc.checkout(force=True)
         self._test_checkout()
 
     def _test_checkout(self):
@@ -44,10 +46,10 @@ class TestCheckout(TestRepro):
 
 class TestCheckoutSingleStage(TestCheckout):
     def test(self):
-        ret = main(['checkout', self.foo_stage.path])
+        ret = main(['checkout', '--force', self.foo_stage.path])
         self.assertEqual(ret, 0)
 
-        ret = main(['checkout', self.data_dir_stage.path])
+        ret = main(['checkout', '--force', self.data_dir_stage.path])
         self.assertEqual(ret, 0)
 
         self._test_checkout()
@@ -62,7 +64,7 @@ class TestCheckoutCorruptedCacheFile(TestRepro):
         with open(cache, 'a') as fd:
             fd.write('1')
 
-        self.dvc.checkout()
+        self.dvc.checkout(force=True)
 
         self.assertFalse(os.path.isfile(self.FOO))
         self.assertFalse(os.path.isfile(cache))
@@ -91,14 +93,14 @@ class TestCheckoutCorruptedCacheDir(TestDvc):
         with open(cache, 'w+') as fobj:
             fobj.write('1')
 
-        self.dvc.checkout()
+        self.dvc.checkout(force=True)
 
         self.assertFalse(os.path.exists(cache))
 
 
 class TestCmdCheckout(TestCheckout):
     def test(self):
-        ret = main(['checkout'])
+        ret = main(['checkout', '--force'])
         self.assertEqual(ret, 0)
         self._test_checkout()
 
@@ -144,7 +146,7 @@ class TestRemoveFilesWhenCheckout(CheckoutBase):
 
         # add the file into a separate branch
         self.dvc.scm.checkout(branch_1, True)
-        ret = main(['checkout'])
+        ret = main(['checkout', '--force'])
         self.assertEqual(ret, 0)
         self.commit_data_file(fname)
 
@@ -154,7 +156,7 @@ class TestRemoveFilesWhenCheckout(CheckoutBase):
 
         # Make sure `dvc checkout` removes the file
         # self.dvc.checkout()
-        ret = main(['checkout'])
+        ret = main(['checkout', '--force'])
         self.assertEqual(ret, 0)
         self.assertFalse(os.path.exists(fname))
 
@@ -162,22 +164,39 @@ class TestRemoveFilesWhenCheckout(CheckoutBase):
 class TestCheckoutCleanWorkingDir(CheckoutBase):
     def test(self):
         from tests.test_data_cloud import sleep
-        
+
         stages = self.dvc.add(self.DATA_DIR)
         stage = stages[0]
 
         sleep()
-        
+
         working_dir_change = os.path.join(self.DATA_DIR, 'not_cached.txt')
         with open(working_dir_change, 'w') as f:
             f.write('not_cached')
 
         sleep()
-            
-        ret = main(['checkout', stage.relpath])
+
+        ret = main(['checkout', '--force', stage.relpath])
         self.assertEqual(ret, 0)
         self.assertFalse(os.path.exists(working_dir_change))
 
+    @patch('dvc.prompt.Prompt.prompt')
+    def test_force(self, mock_prompt):
+        mock_prompt.return_value = False
+
+        stages = self.dvc.add(self.DATA_DIR)
+        stage = stages[0]
+
+        working_dir_change = os.path.join(self.DATA_DIR, 'not_cached.txt')
+        with open(working_dir_change, 'w') as f:
+            f.write('not_cached')
+
+        ret = main(['checkout', stage.relpath])
+
+        mock_prompt.assert_called()
+        self.assertNotEqual(ret, 0)
+        self.assertRaises(DvcException)
+
 
 class TestCheckoutSelectiveRemove(CheckoutBase):
     def test(self):
@@ -193,16 +212,16 @@ class TestCheckoutSelectiveRemove(CheckoutBase):
         staged_files = self.outs_info(stage)
 
         sleep()
-        
+
         os.remove(staged_files[0].path)
-        
+
         sleep()
-        
-        ret = main(['checkout', stage.relpath])
+
+        ret = main(['checkout', '--force', stage.relpath])
         self.assertEqual(ret, 0)
 
         sleep()
-        
+
         checkedout_files = self.outs_info(stage)
 
         self.assertEqual(len(staged_files), len(checkedout_files))
@@ -250,12 +269,12 @@ class TestGitIgnoreWhenCheckout(CheckoutBase):
         self.commit_data_file(fname_master)
 
         self.dvc.scm.checkout(branch_1, True)
-        ret = main(['checkout'])
+        ret = main(['checkout', '--force'])
         self.assertEqual(ret, 0)
         self.commit_data_file(fname_branch)
 
         self.dvc.scm.checkout(branch_master)
-        ret = main(['checkout'])
+        ret = main(['checkout', '--force'])
         self.assertEqual(ret, 0)
 
         ignored = self.read_ignored()
@@ -264,7 +283,7 @@ class TestGitIgnoreWhenCheckout(CheckoutBase):
         self.assertIn('/' + fname_master, ignored)
 
         self.dvc.scm.checkout(branch_1)
-        ret = main(['checkout'])
+        ret = main(['checkout', '--force'])
         self.assertEqual(ret, 0)
         ignored = self.read_ignored()
         self.assertIn('/' + fname_branch, ignored)
@@ -281,7 +300,7 @@ class TestCheckoutMissingMd5InStageFile(TestRepro):
         with open(self.file1_stage, 'w') as fd:
             yaml.dump(d, fd)
 
-        self.dvc.checkout()
+        self.dvc.checkout(force=True)
 
 
 class TestCheckoutEmptyDir(TestDvc):
@@ -298,7 +317,7 @@ class TestCheckoutEmptyDir(TestDvc):
         stage.outs[0].remove()
         self.assertFalse(os.path.exists(dname))
 
-        self.dvc.checkout()
+        self.dvc.checkout(force=True)
 
         self.assertTrue(os.path.isdir(dname))
         self.assertEqual(len(os.listdir(dname)), 0)
@@ -313,7 +332,7 @@ class TestCheckoutNotCachedFile(TestDvc):
                              deps=[self.FOO, self.CODE],
                              outs_no_cache=['out'])
 
-        self.dvc.checkout()
+        self.dvc.checkout(force=True)
 
 
 class TestCheckoutWithDeps(TestRepro):
@@ -324,7 +343,7 @@ class TestCheckoutWithDeps(TestRepro):
         self.assertFalse(os.path.exists(self.FOO))
         self.assertFalse(os.path.exists(self.file1))
 
-        ret = main(['checkout', self.file1_stage, '--with-deps'])
+        ret = main(['checkout', '--force', self.file1_stage, '--with-deps'])
         self.assertEqual(ret, 0)
 
         self.assertTrue(os.path.exists(self.FOO))

tests/test_metrics.py

@@ -214,7 +214,7 @@ class TestNoMetrics(TestDvc):
 class TestCachedMetrics(TestDvc):
     def _do_write(self, branch):
         self.dvc.scm.checkout(branch)
-        self.dvc.checkout()
+        self.dvc.checkout(force=True)
 
         with open('metrics.json', 'w+') as fd:
             json.dump({'metrics': branch}, fd)

tests/test_repro.py

@@ -585,7 +585,7 @@ class TestReproExternalBase(TestDvc):
         self.dvc.gc()
 
         self.dvc.remove(cmd_stage.path, outs_only=True)
-        self.dvc.checkout(cmd_stage.path)
+        self.dvc.checkout(cmd_stage.path, force=True)
 
 
 class TestReproExternalS3(TestReproExternalBase):