From c007f414ea7306514aafc6fb19a87162524b4853 Mon Sep 17 00:00:00 2001
From: Julien Duponchelle <julien@duponchelle.info>
Date: Thu, 28 May 2015 15:05:18 +0200
Subject: [PATCH] Add a progress bar

---
 gns3registry/main.py     | 22 +++++++++++++++++++---
 gns3registry/registry.py | 10 +++++++---
 2 files changed, 26 insertions(+), 6 deletions(-)

diff --git a/gns3registry/main.py b/gns3registry/main.py
index fdf5f4a..f31edbb 100644
--- a/gns3registry/main.py
+++ b/gns3registry/main.py
@@ -19,6 +19,7 @@
 import argparse
 import sys
 import os
+import math
 from distutils.util import strtobool
 
 sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
@@ -26,9 +27,6 @@ sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
 from gns3registry.registry import Registry
 from gns3registry.config import Config, ConfigException
 
-registry = Registry()
-config = Config()
-
 
 def yes_no(message):
     while True:
@@ -37,6 +35,24 @@ def yes_no(message):
         except ValueError:
             pass
 
+def download_progress_callback(count, blockSize, totalSize):
+    """
+    Callback called when a file is downloading
+    """
+
+    if totalSize == -1:
+        sys.stdout.write("Unknow size downloading...\n")
+
+    percent = int(count * blockSize * 100/totalSize)
+    sys.stdout.write("\r[{}{}] {} %".format("#" * math.floor(percent / 2), " " * math.ceil(100 / 2 - percent / 2), percent))
+    if count * blockSize == totalSize:
+        sys.stdout.write("\n")
+    sys.stdout.flush()
+
+registry = Registry(download_progress_callback=download_progress_callback)
+config = Config()
+
+
 def add_images(images):
     print("WARNING WARNING WARNING")
     print("It's experimental")
diff --git a/gns3registry/registry.py b/gns3registry/registry.py
index 4c7f5f0..b49f549 100644
--- a/gns3registry/registry.py
+++ b/gns3registry/registry.py
@@ -27,11 +27,15 @@ from gns3registry.image import Image
 
 
 class Registry:
-    def __init__(self):
-        pass
+    def __init__(self, download_progress_callback=None):
+        """
+        :param download_progress_callback: Callback called when a file is downloaded
+        """
+        self._download_progress_callback = download_progress_callback
 
     def detect_images(self, images_path):
         """
+        :param images_path: Array of path to images
         :returns: Array of configuration corresponding to the image
         """
 
@@ -64,7 +68,7 @@ class Registry:
                         if "direct_download_url" in file:
                             print("Download {} to {}".format(file["direct_download_url"], path))
                             #TODO: Skip download if file already exist with same sha1
-                            urllib.request.urlretrieve(file["direct_download_url"], path)
+                            urllib.request.urlretrieve(file["direct_download_url"], path, reporthook=self._download_progress_callback)
                             return path
                         else:
                             print("You need to manually download the image {filename} from:\n{download_url}\n\nAnd run: ./bin/gns3-get --add {filename}".format(filename=file["filename"], download_url=file["download_url"]))