Source Code for Module TEES.Utils.Download

  1  import sys, os 
  2  import shutil 
  3  import urllib 
  4  import tarfile 
  5  import tempfile 
  6  import zipfile 
  7  from Libraries.progressbar import * 
  8   
  9  pbar = None 
 10   
 11 -def checkReturnCode(code): 
 12      if code != 0: 
 13          print >> sys.stderr, "Non-zero return code", str(code) + ", program may not be working." 
 14          return False 
 15      else: 
 16          print >> sys.stderr, "Return code", str(code) + ", program appears to be working." 
 17          return True 
 18   
 19 -def getTopDir(path, names): 
 20      topDirs = []  
 21      for item in names: 
 22          if "/" not in item or (item.endswith("/") and item.count("/") == 1): 
 23              potential = os.path.join(path, item) 
 24              if os.path.exists(potential) and os.path.isdir(potential):                   
 25                  topDirs.append(item) 
 26      assert len(topDirs) == 1, (path, topDirs, names) 
 27      return os.path.join(path, topDirs[0]) 
 28   
 29  # Modified from http://code.activestate.com/recipes/576714-extract-a-compressed-file/ 
 30 -def extractPackage(path, destPath, subPath=None): 
 31      if path.endswith('.zip'): 
 32          opener, mode = zipfile.ZipFile, 'r' 
 33          namelister = zipfile.ZipFile.namelist 
 34      elif path.endswith('.tar.gz') or path.endswith('.tgz'): 
 35          opener, mode = tarfile.open, 'r:gz' 
 36          namelister = tarfile.TarFile.getnames 
 37      elif path.endswith('.tar.bz2') or path.endswith('.tbz'): 
 38          opener, mode = tarfile.open, 'r:bz2' 
 39          namelister = tarfile.TarFile.getnames 
 40      else:  
 41          raise ValueError, "Could not extract `%s` as no appropriate extractor is found" % path 
 42       
 43      file = opener(path, mode) 
 44      names = namelister(file) 
 45      if subPath == None: 
 46          file.extractall(destPath) 
 47      else: 
 48          tempdir = tempfile.mkdtemp() 
 49          file.extractall(tempdir) 
 50          if os.path.exists(destPath): 
 51              shutil.rmtree(destPath) 
 52          shutil.move(os.path.join(tempdir, subPath), destPath) 
 53          shutil.rmtree(tempdir) 
 54      file.close() 
 55      return names 
 56   
 57 -def downloadAndExtract(url, extractPath=None, downloadPath=None, packagePath=None, addName=True, redownload=False): 
 58      # Download 
 59      downloadFile = download(url, downloadPath, addName=addName, clear=redownload) 
 60      if downloadFile == None: 
 61          return None 
 62      # Unpack 
 63      print >> sys.stderr, "Extracting", downloadFile, "to", extractPath 
 64      return extractPackage(downloadFile, extractPath, packagePath) 
 65   
 66 -def downloadProgress(count, blockSize, totalSize): 
 67      percent = int(count*blockSize*100/totalSize) 
 68      percent = max(0, min(percent, 100)) # clamp 
 69      pbar.update(percent) 
 70   
 71 -def download(url, destPath=None, addName=True, clear=False): 
 72      global pbar 
 73       
 74      redirectedUrl = urllib.urlopen(url).geturl() 
 75      if redirectedUrl != url: 
 76          print >> sys.stderr, "Redirected to", redirectedUrl 
 77      if destPath == None: 
 78          destPath = "/tmp" 
 79      destFileName = destPath 
 80      if addName: 
 81          destFileName = destPath + "/" + os.path.basename(redirectedUrl) 
 82      if not os.path.exists(os.path.dirname(destFileName)): 
 83          os.makedirs(os.path.dirname(destFileName)) 
 84      if clear or not os.path.exists(destFileName): 
 85          if os.path.exists(destFileName): # clear existing file 
 86              os.remove(destFileName) 
 87          print >> sys.stderr, "Downloading file", redirectedUrl, "to", destFileName 
 88          widgets = [FileTransferSpeed(),' <<<', Bar(), '>>> ', Percentage(),' ', ETA()] 
 89          pbar = ProgressBar(widgets=widgets, maxval=100) 
 90          pbar.start() 
 91          try: 
 92              urllib.FancyURLopener().retrieve(redirectedUrl, destFileName, reporthook=downloadProgress) 
 93          except IOError, e: 
 94              print >> sys.stderr, e.errno 
 95              print >> sys.stderr, "Error downloading file", redirectedUrl 
 96              pbar.finish() 
 97              pbar = None 
 98              return None 
 99          pbar.finish() 
100          pbar = None 
101      else: 
102          print >> sys.stderr, "Skipping already downloaded file", url 
103      return destFileName 
104