utils: images downloader: Add support for UHD_IMAGES_URL, clean up

- Remove Python2 compat hacks
- Read the UHD_IMAGES_URL environment variable. If set, it overrides the
  default value for --base-url

1 files changed, 29 insertions, 35 deletions

diff --git a/host/utils/uhd_images_downloader.py.in b/host/utils/uhd_images_downloader.py.in
index 29f559dfe..0e5cc8ed6 100644
--- a/host/utils/uhd_images_downloader.py.in
+++ b/host/utils/uhd_images_downloader.py.in
@@ -45,27 +45,11 @@ import sys
 import tempfile
 import zipfile
 import platform
+from urllib.parse import urljoin  # Python 3
+
 # For all the non-core-library imports, we will be extra paranoid and be very
 # nice with error messages so that everyone understands what's up.
 try:
-    from urllib.parse import urljoin  # Python 3
-except ImportError:
-    from urlparse import urljoin      # Python 2
-try:
-    from builtins import input
-except ImportError:
-    input = raw_input
-try:
-    from six import iteritems
-except ImportError:
-    sys.stdout.write(
-        "[ERROR] Missing module 'six'! Please install it, e.g., by "
-        "running 'pip install six' or any other tool that can install "
-        "Python modules.\n")
-    if platform.system() == 'Windows':
-        input('Hit Enter to continue.')
-    exit(0)
-try:
     import requests
 except ImportError:
     sys.stdout.write(
@@ -76,8 +60,7 @@ except ImportError:
         input('Hit Enter to continue.')
     exit(0)
 
-
-
+# pylint: disable=bad-whitespace
 _DEFAULT_TARGET_REGEX     = "(fpga|fw|windrv)_default"
 _BASE_DIR_STRUCTURE_PARTS = ["share", "uhd", "images"]
 _DEFAULT_INSTALL_PATH     = os.path.join("@CMAKE_INSTALL_PREFIX@", *_BASE_DIR_STRUCTURE_PARTS)
@@ -88,16 +71,19 @@ _DEFAULT_BUFFER_SIZE      = 8192
 _DEFAULT_DOWNLOAD_LIMIT   = 100 * 1024 * 1024 # Bytes
 _ARCHIVE_DEFAULT_TYPE     = "zip"
 _UHD_VERSION              = "@UHD_VERSION@"
-# Note: _MANIFEST_CONTENTS are placed at the bottom of this file for aesthetic reasons
-_LOG_LEVELS = {"TRACE": 1,
-               "DEBUG": 2,
-               "INFO": 3,
-               "WARN": 4,
-               "ERROR": 5}
+_LOG_LEVELS = {
+    "TRACE": 1,
+    "DEBUG": 2,
+    "INFO":  3,
+    "WARN":  4,
+    "ERROR": 5,
+}
 _LOG_LEVEL = _LOG_LEVELS["INFO"]
 _YES = False
 _PROXIES = {}
-
+_BASE_URL_ENV_VAR = "UHD_IMAGES_URL"
+# pylint: enable=bad-whitespace
+# Note: _MANIFEST_CONTENTS are placed at the bottom of this file for aesthetic reasons
 
 def log(level, message):
     """Logging function"""
@@ -139,8 +125,10 @@ def parse_args():
                         "http://user:pass@1.2.3.4:port\n"
                         "If this this option is not given, the environment "
                         "variable HTTP_PROXY can also be used to specify a proxy.")
-    parser.add_argument("-b", "--base-url", type=str, default=_DEFAULT_BASE_URL,
-                        help="Set base URL for images download location")
+    parser.add_argument("-b", "--base-url", type=str,
+                        help="Set base URL for images download location. "
+                        "Defaults to ${} if set, or {} otherwise.".format(
+                            _BASE_URL_ENV_VAR, _DEFAULT_BASE_URL))
     parser.add_argument("-k", "--keep", action="store_true", default=False,
                         help="Keep the downloaded images archives in the image directory")
     parser.add_argument("-T", "--test", action="store_true", default=False,
@@ -156,8 +144,18 @@ def parse_args():
                         help="Decrease verbosity level")
     parser.add_argument('-v', '--verbose', action='count', default=0,
                         help="Increase verbosity level")
-    # Some sanitation that's easier to handle outside of the argparse framework:
     args = parser.parse_args()
+    # Set the verbosity
+    global _LOG_LEVEL
+    log("TRACE", "Default log level: {}".format(_LOG_LEVEL))
+    _LOG_LEVEL = _LOG_LEVEL - args.verbose + args.quiet
+    # Some sanitation that's easier to handle outside of the argparse framework:
+    if not args.base_url:
+        if os.environ.get(_BASE_URL_ENV_VAR):
+            args.base_url = os.environ.get(_BASE_URL_ENV_VAR)
+        else:
+            args.base_url = _DEFAULT_BASE_URL
+    log("INFO", "Using base URL: {}".format(args.base_url))
     if not args.base_url.endswith('/') and args.base_url != "":
         args.base_url += '/'
     if args.yes:
@@ -166,10 +164,6 @@ def parse_args():
     if args.http_proxy:
         global _PROXIES
         _PROXIES['http'] = args.http_proxy
-    # Set the verbosity
-    global _LOG_LEVEL
-    log("TRACE", "Default log level: {}".format(_LOG_LEVEL))
-    _LOG_LEVEL = _LOG_LEVEL - args.verbose + args.quiet
     return args
 
 
@@ -343,7 +337,7 @@ def print_target_list(manifest, args):
                 base=args.base_url,
                 relpath=value["url"]))
     else:
-        for manifest_item in iteritems(manifest):
+        for manifest_item in manifest.items():
             print(args.base_url+manifest_item[1]["url"])