Implement support for fsspec callbacks in put_file/get_file

adlfs/spec.py

@@ -287,6 +287,25 @@ def seek(self, loc: int, whence: int = 0, **kwargs):
         return self.loc
 
 
+# https://github.com/Azure/azure-sdk-for-python/issues/11419#issuecomment-628143480
+def make_callback(key, callback):
+    if callback is None:
+        return None
+
+    sent_total = False
+
+    def wrapper(response):
+        nonlocal sent_total
+
+        current = response.context.get(key)
+        total = response.context["data_stream_total"]
+        if not sent_total:
+            callback.set_size(total)
+        callback.absolute_update(current)
+
+    return wrapper
+
+
 class AzureBlobFileSystem(AsyncFileSystem):
     """
     Access Azure Datalake Gen2 and Azure Storage if it were a file system using Multiprotocol Access
@@ -1525,7 +1544,9 @@ async def _expand_path(self, path, recursive=False, maxdepth=None, **kwargs):
             raise FileNotFoundError
         return list(sorted(out))
 
-    async def _put_file(self, lpath, rpath, delimiter="/", overwrite=False, **kwargws):
+    async def _put_file(
+        self, lpath, rpath, delimiter="/", overwrite=False, callback=None, **kwargws
+    ):
         """
         Copy single file to remote
 
@@ -1546,7 +1567,12 @@ async def _put_file(self, lpath, rpath, delimiter="/", overwrite=False, **kwargw
                         container_name, path
                     ) as bc:
                         await bc.upload_blob(
-                            f1, overwrite=overwrite, metadata={"is_directory": "false"}
+                            f1,
+                            overwrite=overwrite,
+                            metadata={"is_directory": "false"},
+                            raw_response_hook=make_callback(
+                                "upload_stream_current", callback
+                            ),
                         )
                 self.invalidate_cache()
             except ResourceExistsError:
@@ -1585,24 +1611,24 @@ def download(self, rpath, lpath, recursive=False, **kwargs):
         """Alias of :ref:`FilesystemSpec.get`."""
         return self.get(rpath, lpath, recursive=recursive, **kwargs)
 
-    async def _get_file(self, rpath, lpath, recursive=False, delimiter="/", **kwargs):
+    async def _get_file(
+        self, rpath, lpath, recursive=False, delimiter="/", callback=None, **kwargs
+    ):
         """ Copy single file remote to local """
-        files = await self._ls(rpath)
-        files = [f["name"] for f in files]
         container_name, path = self.split_path(rpath, delimiter=delimiter)
         try:
-            if await self._isdir(rpath):
-                os.makedirs(lpath, exist_ok=True)
-            else:
-                async with self.service_client.get_blob_client(
-                    container_name, path.rstrip(delimiter)
-                ) as bc:
-                    with open(lpath, "wb") as my_blob:
-                        stream = await bc.download_blob()
-                        data = await stream.readall()
-                        my_blob.write(data)
-        except Exception as e:
-            raise FileNotFoundError(f"File not found for {e}")
+            async with self.service_client.get_blob_client(
+                container_name, path.rstrip(delimiter)
+            ) as bc:
+                with open(lpath, "wb") as my_blob:
+                    stream = await bc.download_blob(
+                        raw_response_hook=make_callback(
+                            "download_stream_current", callback
+                        )
+                    )
+                    await stream.readinto(my_blob)
+        except ResourceNotFoundError as exception:
+            raise FileNotFoundError from exception
 
     get_file = sync_wrapper(_get_file)