fs.get_file: add callback support (#137)

Author: skshetry

pydrive2/fs/spec.py

@@ -364,21 +364,30 @@ def cp_file(self, lpath, rpath, **kwargs):
             buffer = io.BytesIO(stream.read())
             self.upload_fobj(buffer, rpath)
 
-    def get_file(self, lpath, rpath, callback=None, **kwargs):
+    def get_file(self, lpath, rpath, callback=None, block_size=None, **kwargs):
         item_id = self._get_item_id(lpath)
-        return self.gdrive_get_file(item_id, rpath, callback)
+        return self.gdrive_get_file(
+            item_id, rpath, callback=callback, block_size=block_size
+        )
 
     @_gdrive_retry
-    def gdrive_get_file(self, item_id, rpath, callback):
+    def gdrive_get_file(self, item_id, rpath, callback=None, block_size=None):
         param = {"id": item_id}
         # it does not create a file on the remote
         gdrive_file = self.client.CreateFile(param)
 
         extra_args = {}
+        if block_size:
+            extra_args["chunksize"] = block_size
+
         if callback:
+
+            def cb(value, _):
+                callback.absolute_update(value)
+
             gdrive_file.FetchMetadata(fields="fileSize")
             callback.set_size(int(gdrive_file.get("fileSize")))
-            extra_args["callback"] = callback.update
+            extra_args["callback"] = cb
 
         gdrive_file.GetContentFile(rpath, **extra_args)
 

pydrive2/test/test_fs.py

@@ -4,7 +4,7 @@
 from concurrent import futures
 
 import pytest
-
+import fsspec
 from pydrive2.auth import GoogleAuth
 from pydrive2.fs import GDriveFileSystem
 from pydrive2.test.test_util import settings_file_path, setup_credentials
@@ -210,3 +210,21 @@ def test_get_file(fs, tmpdir, remote_dir):
     fs.put_file(src_file, remote_dir + "/a.txt")
     fs.get_file(remote_dir + "/a.txt", dest_file)
     assert dest_file.read() == "data"
+
+
+def test_get_file_callback(fs, tmpdir, remote_dir):
+    src_file = tmpdir / "a.txt"
+    dest_file = tmpdir / "b.txt"
+
+    with open(src_file, "wb") as file:
+        file.write(b"data" * 10)
+
+    fs.put_file(src_file, remote_dir + "/a.txt")
+    callback = fsspec.Callback()
+    fs.get_file(
+        remote_dir + "/a.txt", dest_file, callback=callback, block_size=10
+    )
+    assert dest_file.read() == "data" * 10
+
+    assert callback.size == 40
+    assert callback.value == 40