Commit e1f0fe3c authored by Martin Braun's avatar Martin Braun Committed by michael-west

utils: uhd_images_downloader: Add --yes option

This will answer 'yes' to all questions that are asked and allows to
script the downloader. Example:

    $ uhd_images_downloader -t sdimg -y

Any interaction that the script will have with the user will be skipped.
parent 39cc1d8c
......@@ -22,6 +22,10 @@ try:
from urllib.parse import urljoin # Python 3
except ImportError:
from urlparse import urljoin # Python 2
from builtins import input
except ImportError:
input = raw_input
from six import iteritems
import requests
......@@ -43,6 +47,7 @@ _LOG_LEVELS = {"TRACE": 1,
"WARN": 4,
"ERROR": 5}
_YES = False
def log(level, message):
......@@ -79,6 +84,10 @@ def parse_args():
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,
help="Verify the downloaded archives before extracting them")
parser.add_argument("-y", "--yes", action="store_true", default=False,
help="Answer all questions with 'yes' (for scripting purposes).")
parser.add_argument("-n", "--dry-run", action="store_true", default=False,
help="Print selected target without actually downloading them.")
parser.add_argument("--refetch", action="store_true", default=False,
......@@ -91,6 +100,24 @@ def parse_args():
return parser.parse_args()
def ask_permission(question, default_no=True):
Ask the question, and have the user type y or n on the keyboard. If the
global variable _YES is true, this always returns True without asking the
question. Otherwise, return True if the answer is 'yes', 'y', or something
if _YES:
log("DEBUG", "Assuming the answer is 'yes' for this question: " +
return True
postfix = "[y/N]" if default_no else "[Y/n]"
answer = input(question + " " + postfix)
if answer and answer[0].lower() == 'y':
return True
return False
class TemporaryDirectory:
"""Class to create a temporary directory"""
def __enter__(self):
......@@ -266,11 +293,22 @@ def delete_from_inv(target_info, inventory, images_dir):
return True
def extract(archive_path, images_dir, archive_type):
def extract(archive_path, images_dir, archive_type, test_zip=False):
"""Extract the contents of the archive into `images_dir`"""
if archive_type == "zip":
log("TRACE", "Attempting to extracted files from {}".format(archive_path))
with zipfile.ZipFile(archive_path) as images_zip:
# Check that the Zip file is valid, in which case `testzip()` returns None.
# If its bad, that function will return a list of bad files
if test_zip and images_zip.testzip():
log("ERROR", "Could not extract the following invalid Zip file:"
" {}".format(archive_path))
return []
except OSError:
log("ERROR", "Could not extract the following invalid Zip file:"
" {}".format(archive_path))
return []
archive_namelist = images_zip.namelist()
log("TRACE", "Extracted files: {}".format(archive_namelist))
......@@ -400,7 +438,7 @@ def main():
# Otherwise, the check has succeeded, and we can proceed
delete_from_inv(target_info, inventory, images_dir)
archive_namelist = extract(temp_path, images_dir, archive_type)
archive_namelist = extract(temp_path, images_dir, archive_type, args.test)
if args.keep:
# If the user wants to keep the downloaded archive,
# save it to the images directory and add it to the inventory
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment