Added a callback to show file transfer progress (#78)

* Added a callback to show file transfer progress

* Avoid notifying of progress again when data was fully written

* Wrap progress_callback with try/catch

* Separate progress callback code

* Make progress generator reusable

* Implemented pull progress callback

Author: MohammadAG

adb/adb_commands.py

@@ -109,7 +109,7 @@ def Devices(cls):
   def GetState(self):
     return self._device_state
 
-  def Install(self, apk_path, destination_dir='', replace_existing=True, timeout_ms=None):
+  def Install(self, apk_path, destination_dir='', timeout_ms=None, replace_existing=True, transfer_progress_callback=None):
     """Install an apk to the device.
 
     Doesn't support verifier file, instead allows destination directory to be
@@ -121,6 +121,7 @@ def Install(self, apk_path, destination_dir='', replace_existing=True, timeout_m
         persistent applications.
       replace_existing: whether to replace existing application
       timeout_ms: Expected timeout for pushing and installing.
+      transfer_progress_callback: callback method that accepts filename, bytes_written and total_bytes of APK transfer
 
     Returns:
       The pm install output.
@@ -129,15 +130,14 @@ def Install(self, apk_path, destination_dir='', replace_existing=True, timeout_m
       destination_dir = '/data/local/tmp/'
     basename = os.path.basename(apk_path)
     destination_path = destination_dir + basename
-    self.Push(apk_path, destination_path, timeout_ms=timeout_ms)
-
+    self.Push(apk_path, destination_path, timeout_ms=timeout_ms, progress_callback=transfer_progress_callback)
+    
     cmd = ['pm install']
     if replace_existing:
       cmd.append('-r')
     cmd.append('"%s"' % destination_path)
-
     return self.Shell(' '.join(cmd), timeout_ms=timeout_ms)
-
+  
   def Uninstall(self, package_name, keep_data=False, timeout_ms=None):
     """Removes a package from the device.
 
@@ -155,7 +155,7 @@ def Uninstall(self, package_name, keep_data=False, timeout_ms=None):
     cmd.append('"%s"' % package_name)
     return self.Shell(' '.join(cmd), timeout_ms=timeout_ms)
 
-  def Push(self, source_file, device_filename, mtime='0', timeout_ms=None):
+  def Push(self, source_file, device_filename, mtime='0', timeout_ms=None, progress_callback=None):
     """Push a file or directory to the device.
 
     Args:
@@ -164,26 +164,28 @@ def Push(self, source_file, device_filename, mtime='0', timeout_ms=None):
       device_filename: Destination on the device to write to.
       mtime: Optional, modification time to set on the file.
       timeout_ms: Expected timeout for any part of the push.
+      progress_callback: callback method that accepts filename, bytes_written and total_bytes,
+                         total_bytes will be -1 for file-like objects
     """
     should_close = False
     if isinstance(source_file, str):
       if os.path.isdir(source_file):
         self.Shell("mkdir " + device_filename)
         for f in os.listdir(source_file):
-          self.Push(os.path.join(source_file, f), device_filename + '/' + f)
+          self.Push(os.path.join(source_file, f), device_filename + '/' + f, progress_callback=progress_callback)
         return
       source_file = open(source_file, "rb")
       should_close = True
 
     connection = self.protocol_handler.Open(
         self.handle, destination=b'sync:', timeout_ms=timeout_ms)
     self.filesync_handler.Push(connection, source_file, device_filename,
-                               mtime=int(mtime))
+                               mtime=int(mtime), progress_callback=progress_callback)
     if should_close:
       source_file.close()
     connection.Close()
 
-  def Pull(self, device_filename, dest_file='', timeout_ms=None):
+  def Pull(self, device_filename, dest_file='', progress_callback=None, timeout_ms=None):
     """Pull a file from the device.
 
     Args:
@@ -201,7 +203,7 @@ def Pull(self, device_filename, dest_file='', timeout_ms=None):
     connection = self.protocol_handler.Open(
         self.handle, destination=b'sync:',
         timeout_ms=timeout_ms)
-    self.filesync_handler.Pull(connection, device_filename, dest_file)
+    self.filesync_handler.Pull(connection, device_filename, dest_file, progress_callback)
     connection.Close()
     if isinstance(dest_file, io.BytesIO):
       return dest_file.getvalue()

adb/filesync_protocol.py

@@ -18,6 +18,7 @@
 """
 
 import collections
+import os
 import stat
 import struct
 import time
@@ -76,18 +77,36 @@ def List(cls, connection, path):
     return files
 
   @classmethod
-  def Pull(cls, connection, filename, dest_file):
+  def Pull(cls, connection, filename, dest_file, progress_callback):
     """Pull a file from the device into the file-like dest_file."""
+    if progress_callback:
+      total_bytes = cls.Stat(connection, filename)[1]
+      progress = cls._HandleProgress(lambda current: progress_callback(filename, current, total_bytes))
+      next(progress)
+
     cnxn = FileSyncConnection(connection, b'<2I')
     cnxn.Send(b'RECV', filename)
     for cmd_id, _, data in cnxn.ReadUntil((b'DATA',), b'DONE'):
       if cmd_id == b'DONE':
         break
       dest_file.write(data)
+      if progress_callback:
+        progress.send(len(data))
+
+  @classmethod
+  def _HandleProgress(cls, progress_callback):
+    """Calls the callback with the current progress and total ."""
+    current = 0
+    while True:
+      current += yield
+      try:
+        progress_callback(current)
+      except Exception:  # pylint: disable=broad-except
+        continue
 
   @classmethod
   def Push(cls, connection, datafile, filename,
-           st_mode=DEFAULT_PUSH_MODE, mtime=0):
+           st_mode=DEFAULT_PUSH_MODE, mtime=0, progress_callback=None):
     """Push a file-like object to the device.
 
     Args:
@@ -96,6 +115,7 @@ def Push(cls, connection, datafile, filename,
       filename: Filename to push to
       st_mode: stat mode for filename
       mtime: modification time
+      progress_callback: callback method that accepts filename, bytes_written and total_bytes,
 
     Raises:
       PushFailedError: Raised on push failure.
@@ -107,11 +127,20 @@ def Push(cls, connection, datafile, filename,
     cnxn = FileSyncConnection(connection, b'<2I')
     cnxn.Send(b'SEND', fileinfo)
 
+    if progress_callback:
+      total_bytes = os.fstat(datafile.fileno()).st_size if isinstance(datafile, file) else -1
+      progress = cls._HandleProgress(lambda current: progress_callback(filename, current, total_bytes))
+      next(progress)
+
     while True:
       data = datafile.read(MAX_PUSH_DATA)
-      if not data:
+      if data:
+        cnxn.Send(b'DATA', data)
+
+        if progress_callback:
+          progress.send(len(data))
+      else:
         break
-      cnxn.Send(b'DATA', data)
 
     if mtime == 0:
       mtime = int(time.time())