uhd_images_downloader: more improvements/fixes

* Better error handling
* Improved default download location logic

1 files changed, 24 insertions, 14 deletions

diff --git a/host/utils/uhd_images_downloader.py.in b/host/utils/uhd_images_downloader.py.in
index 87c148258..91f32825f 100644
--- a/host/utils/uhd_images_downloader.py.in
+++ b/host/utils/uhd_images_downloader.py.in
@@ -45,16 +45,23 @@ class temp_dir():
         return self.name
     def __exit__(self, type, value, traceback):
         try:
-            os.removedirs(self.name)
-        except:
+            shutil.rmtree(self.name)
+        except OSError:
             #Utility should have already detected this, but this is for safety
             raise Exception("Could not install images! Make sure you have write permissions.")
 
 if __name__ == "__main__":
 
+    print
+    if os.environ.get("UHD_IMAGES_DIR") != None and os.environ.get("UHD_IMAGES_DIR") != "":
+        default_images_dir = os.environ.get("UHD_IMAGES_DIR")
+        print "UHD_IMAGES_DIR environment variable is set. Default install location: %s" % default_images_dir
+    else:
+        default_images_dir = "@CMAKE_INSTALL_PREFIX@/share/uhd/images"
+
     #Command line options
     parser = OptionParser()
-    parser.add_option("--install-location", type="string", default="", help="Set custom install location for images")
+    parser.add_option("--install-location", type="string", default=default_images_dir, help="Set custom install location for images")
     parser.add_option("--buffer-size", type="int", default=8192, help="Set download buffer size, [default=%default]",)
     (options, args) = parser.parse_args()
     
@@ -69,18 +76,21 @@ if __name__ == "__main__":
     with temp_dir() as dirname:
         os.chdir(dirname)
 
-        #Configuring image destination
-        if options.install_location != "":
-            images_dir = os.path.abspath(os.path.join(current_directory, options.install_location))
-        elif os.environ.get("UHD_IMAGES_DIR") != "" and os.environ.get("UHD_IMAGES_DIR") != None:
-            images_dir = os.environ.get("UHD_IMAGES_DIR")
+        if options.install_location == default_images_dir:
+            images_dir = default_images_dir
+        elif os.path.isabs(options.install_location):
+            #Custom absolute path given
+            images_dir = options.install_location
         else:
-            images_dir = "@CMAKE_INSTALL_PREFIX@/share/uhd/images"
-        
+            #Custom relative path given, so construct absolute path
+            images_dir = os.path.abspath(os.path.join(current_directory, options.install_location))
+
         #Before doing anything, check for write permissions in parent directory
         parent_directory = os.path.dirname(images_dir)
-        if not os.access(parent_directory, os.W_OK):
-            sys.stderr.write("You do not have write permissions at the install location!\n")
+        if os.access(parent_directory, os.W_OK):
+            print "Downloading images to: %s" % images_dir
+        else:
+            print "You do not have write permissions at the install location!"
             sys.exit(1)
 
         opener = urllib2.build_opener()
@@ -129,8 +139,8 @@ if __name__ == "__main__":
             if os.path.exists(images_dir):
                 try:
                     shutil.rmtree(images_dir)
-                except:
-                    sys.stderr.write("\nMake sure you have write permissions in the images directory.\n")
+                except OSError:
+                    print "Make sure you have write permissions in the images directory."
                     sys.exit(1)
 
             #Copying downloaded images into images_dir