Repository: iterative/dvc
Issue: Number of already present files should be reflected in push progress

Hey guys, IMHO the number of pushed files should also include the files already found in the remote cache.

For example, when I upload 2 new files out of my cache of 1000 files, the progress should say 1000/1000 at the end, instead of 2/1000.

Changed files:
dvc/cloud/data_cloud.py, dvc/cloud/base.py

diff --git a/dvc/cloud/base.py b/dvc/cloud/base.py
index 54dec99dd..21346e654 100644
--- a/dvc/cloud/base.py
+++ b/dvc/cloud/base.py
@@ -221,3 +221,9 @@ class DataCloudBase(object):
 
     def disconnect(self):
         pass
+
+    def __enter__(self):
+        self.connect()
+
+    def __exit__(self, type, value, tb):
+        self.disconnect()
diff --git a/dvc/cloud/data_cloud.py b/dvc/cloud/data_cloud.py
index 06b23be65..9949e1b1f 100644
--- a/dvc/cloud/data_cloud.py
+++ b/dvc/cloud/data_cloud.py
@@ -11,6 +11,7 @@ from dvc.cloud.gcp import DataCloudGCP
 from dvc.cloud.ssh import DataCloudSSH
 from dvc.cloud.local import DataCloudLOCAL
 from dvc.cloud.base import DataCloudBase, CloudSettings
+from dvc.cloud.base import STATUS_MODIFIED, STATUS_NEW, STATUS_DELETED
 
 
 class DataCloud(object):
@@ -99,47 +100,54 @@ class DataCloud(object):
 
         return collected
 
-    def _map_targets(self, func, targets, jobs, collect_local=False, collect_cloud=False, remote=None):
-        """
-        Process targets as data items in parallel.
-        """
+    def _get_cloud(self, remote):
+        if remote:
+            return self._init_remote(remote)
 
-        if not remote:
-            cloud = self._cloud
-        else:
-            cloud = self._init_remote(remote)
+        return self._cloud
 
+    def _filter(self, func, status, targets, jobs, remote):
+        cloud = self._get_cloud(remote)
         if not cloud:
-            return
-
-        cloud.connect()
-
-        collected = set()
-        if collect_local:
-            collected |= self._collect(cloud, targets, jobs, True)
-        if collect_cloud:
-            collected |= self._collect(cloud, targets, jobs, False)
+            return []
 
-        ret = map_progress(getattr(cloud, func), list(collected), jobs)
+        with cloud:
+            filtered = []
+            for t, s in self._status(cloud, targets, jobs):
+                if s == STATUS_MODIFIED or s == status:
+                    filtered.append(t)
 
-        cloud.disconnect()
-
-        return ret
+            return map_progress(getattr(cloud, func), filtered, jobs)
 
     def push(self, targets, jobs=1, remote=None):
         """
         Push data items in a cloud-agnostic way.
         """
-        return self._map_targets('push', targets, jobs, collect_local=True, remote=remote)
+        return self._filter('push', STATUS_NEW, targets, jobs, remote)
 
     def pull(self, targets, jobs=1, remote=None):
         """
         Pull data items in a cloud-agnostic way.
         """
-        return self._map_targets('pull', targets, jobs, collect_cloud=True, remote=remote)
+        return self._filter('pull', STATUS_DELETED, targets, jobs, remote)
+
+    def _collect_targets(self, cloud, targets, jobs=1):
+        collected = set()
+        collected |= self._collect(cloud, targets, jobs, True)
+        collected |= self._collect(cloud, targets, jobs, False)
+        return list(collected)
+
+    def _status(self, cloud, targets, jobs=1):
+        collected = self._collect_targets(cloud, targets, jobs)
+        return map_progress(cloud.status, collected, jobs)
 
     def status(self, targets, jobs=1, remote=None):
         """
         Check status of data items in a cloud-agnostic way.
         """
-        return self._map_targets('status', targets, jobs, True, True, remote=remote)
+        cloud = self._get_cloud(remote)
+        if not cloud:
+            return []
+
+        with cloud:
+            return self._status(cloud, targets, jobs)