From 5b8937d0fdbd6444c0398ae277c10d28112d692e Mon Sep 17 00:00:00 2001
From: David Young <david@textflex.com>
Date: Tue, 12 Jul 2022 11:28:50 +0800
Subject: [PATCH 1/2] Add an optional attribute to handle updates during atlas
 downloads

Atlas download progress is currently handled within a Rich progress indicator, but external libraries do not have access to these updates. Add an optional attribute to `BrainGlobeAtlas` to reference a handler that is updated during atlas downloads. This handler is a function that takes the currently competed and total required bytes to download the given atlas.
---
 bg_atlasapi/bg_atlas.py |  7 ++++++-
 bg_atlasapi/utils.py    | 19 ++++++++++++++++---
 2 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/bg_atlasapi/bg_atlas.py b/bg_atlasapi/bg_atlas.py
index 33eca11e..dd81d02f 100644
--- a/bg_atlasapi/bg_atlas.py
+++ b/bg_atlasapi/bg_atlas.py
@@ -38,6 +38,9 @@ class BrainGlobeAtlas(core.Atlas):
         instantiation and to suppress warnings.
     print_authors : bool (optional)
         If true, disable default listing of the atlas reference.
+    fn_update : Callable
+        Handler function to update during download. Takes completed and total
+        bytes.
 
     """
 
@@ -51,8 +54,10 @@ def __init__(
         interm_download_dir=None,
         check_latest=True,
         config_dir=None,
+        fn_update=None,
     ):
         self.atlas_name = atlas_name
+        self.fn_update = fn_update
 
         # Read BrainGlobe configuration file:
         conf = config.read_config(config_dir)
@@ -156,7 +161,7 @@ def download_extract_file(self):
         destination_path = self.interm_download_dir / COMPRESSED_FILENAME
 
         # Try to download atlas data
-        utils.retrieve_over_http(self.remote_url, destination_path)
+        utils.retrieve_over_http(self.remote_url, destination_path, self.fn_update)
 
         # Uncompress in brainglobe path:
         tar = tarfile.open(destination_path)
diff --git a/bg_atlasapi/utils.py b/bg_atlasapi/utils.py
index bd76fc17..a7d7fef2 100644
--- a/bg_atlasapi/utils.py
+++ b/bg_atlasapi/utils.py
@@ -1,6 +1,7 @@
 import configparser
 import json
 import logging
+from typing import Callable
 
 import requests
 import tifffile
@@ -128,7 +129,8 @@ def check_internet_connection(
     return False
 
 
-def retrieve_over_http(url, output_file_path):
+def retrieve_over_http(url, output_file_path,
+                       fn_update: Callable[[int, int], None] = None):
     """Download file from remote location, with progress bar.
 
     Parameters
@@ -137,6 +139,9 @@ def retrieve_over_http(url, output_file_path):
         Remote URL.
     output_file_path : str or Path
         Full file destination for download.
+    fn_update : Callable
+        Handler function to update during download. Takes completed and total
+        bytes.
 
     """
     # Make Rich progress bar
@@ -157,17 +162,25 @@ def retrieve_over_http(url, output_file_path):
 
     try:
         with progress:
+            tot = int(response.headers.get("content-length", 0))
             task_id = progress.add_task(
                 "download",
                 filename=output_file_path.name,
                 start=True,
-                total=int(response.headers.get("content-length", 0)),
+                total=tot,
             )
 
             with open(output_file_path, "wb") as fout:
+                advanced = 0
                 for chunk in response.iter_content(chunk_size=CHUNK_SIZE):
                     fout.write(chunk)
-                    progress.update(task_id, advance=len(chunk), refresh=True)
+                    adv = len(chunk)
+                    progress.update(task_id, advance=adv, refresh=True)
+                    
+                    if fn_update:
+                        # update handler with completed and total bytes
+                        advanced += adv
+                        fn_update(advanced, tot)
 
     except requests.exceptions.ConnectionError:
         output_file_path.unlink()

From 18514711db58d9ccd29b6f9b129cc989701f7974 Mon Sep 17 00:00:00 2001
From: David Young <david@textflex.com>
Date: Thu, 14 Jul 2022 10:46:45 +0800
Subject: [PATCH 2/2] Fix formatting errors picked up by black package

- Remove whitespace on empty line
- Format line width, indentation, and closing parenthesis placement
---
 bg_atlasapi/bg_atlas.py | 4 +++-
 bg_atlasapi/utils.py    | 7 ++++---
 2 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/bg_atlasapi/bg_atlas.py b/bg_atlasapi/bg_atlas.py
index dd81d02f..13200e63 100644
--- a/bg_atlasapi/bg_atlas.py
+++ b/bg_atlasapi/bg_atlas.py
@@ -161,7 +161,9 @@ def download_extract_file(self):
         destination_path = self.interm_download_dir / COMPRESSED_FILENAME
 
         # Try to download atlas data
-        utils.retrieve_over_http(self.remote_url, destination_path, self.fn_update)
+        utils.retrieve_over_http(
+            self.remote_url, destination_path, self.fn_update
+        )
 
         # Uncompress in brainglobe path:
         tar = tarfile.open(destination_path)
diff --git a/bg_atlasapi/utils.py b/bg_atlasapi/utils.py
index a7d7fef2..6b10b62d 100644
--- a/bg_atlasapi/utils.py
+++ b/bg_atlasapi/utils.py
@@ -129,8 +129,9 @@ def check_internet_connection(
     return False
 
 
-def retrieve_over_http(url, output_file_path,
-                       fn_update: Callable[[int, int], None] = None):
+def retrieve_over_http(
+    url, output_file_path, fn_update: Callable[[int, int], None] = None
+):
     """Download file from remote location, with progress bar.
 
     Parameters
@@ -176,7 +177,7 @@ def retrieve_over_http(url, output_file_path,
                     fout.write(chunk)
                     adv = len(chunk)
                     progress.update(task_id, advance=adv, refresh=True)
-                    
+
                     if fn_update:
                         # update handler with completed and total bytes
                         advanced += adv