# Copyright (c) 2010 ActiveState Software Inc. All rights reserved. """ net.py ~~~~~~ Contains a single `download_file` method to download from the network with support for, 1. progress bar (via applib.textui) support. 2. ETag 3. download cache """ from __future__ import unicode_literals import os.path as P import logging import json from datetime import datetime from collections import namedtuple from applib import sh from applib import textui import six import six.moves if six.PY3: from base64 import encodebytes else: from base64 import encodestring as encodebytes from pypm.common.util import BareDateTime from pypm.common.util import get_user_agent __all__ = ['download_file', 'DownloadError'] LOG = logging.getLogger(__name__) def download_file(url, target_dir, urlgetter_kargs=None, interactive=True): """Download ``url`` to ``target_dir`` - urlgetter_kargs: kwargs to URLGetter.get(...) method - interactive: if True, show progress bar. Raise ``DownloadError`` when there is an error Return (downloaded_file, actually_downloaded?) """ assert '\n' not in url, 'url has whitespace' url = url.split('#', 1)[0] # remove url anchor if not urlgetter_kargs: urlgetter_kargs = {} # basename apparently works on URLs too # http://stackoverflow.com/questions/1112545/ target_file = P.abspath(P.join(target_dir, P.basename(url))) LOG.debug('Downloading %s', url) if interactive: dl = DownloaderWithProgress(url) else: dl = Downloader(url) downloaded = dl.get(target_file, **urlgetter_kargs) return target_file, downloaded # raised by `Downloader.get` class DownloadError(Exception): """Wrapping exception for all network related errors""" # -- internal support code class URLGetter(object): """A fine-grained URL getter. An improved `urllib.urlretrieve`""" def __init__(self, url): """ - url: URL to download """ self.url = url def get(self, target_file, start_info=None, use_cache=False, save_properties=False, auth=None): """Download the URL and save it to `target_file` Appropriate hook functions are automatically invoked - start_info: message to print (LOG.info) when transfer begins, cache is validated or or an error occurs; it can also be a callable. - use_cache: if True, download will not actually happen if `target_file` is already available *and* its ETag is not changed in the server. For this to work, the previous call to `get` must have passed `save_properties=True` (in order to save the ETag). - save_properties: if True, a file named .$target_file.urlprops containing request headers and other metadata will be saved along side the target_file. (This is required for `use_cache` in future downloads) - auth: (username, password) -- optional http basic auth data Return True only if the download actually happened. """ def invoke_start_info(status): if not start_info: return if six.callable(start_info): i = start_info(status) else: i = start_info.format(status=status) if i: LOG.info(i) if not P.exists(target_file): use_cache = False urlprops = URLProperties(target_file) props = urlprops.load() if props: # write back the new value of `last_attempt_utc` *now* so we don't # have to deal with it when an exception arises later. # last_attempt_utc is simpy the time of the last download attempt props.custom['last_attempt_utc'] = BareDateTime.to_string(datetime.utcnow()) urlprops.save(props.headers, props.custom) req = None if use_cache and props: # Enable the cache header `If-None-Match` etag = props.headers.get('ETag', props.headers.get('etag', None)) if etag: req = six.moves.Request(self.url, headers={'If-None-Match': etag}) else: LOG.warn('no ETag in last headers: %s', props.headers) req = six.moves.Request(self.url) else: req = six.moves.Request(self.url) if auth: username, password = auth req.add_header('Authorization', _create_http_basicauth_header( username, password )) # Set User-Agent # XXX: (in 2.6) urllib2.py does not expose its default user-agent string # so we copy-paste that code here (from urllib2.OpenerDirector) urllib2_user_agent = "Python-urllib/%s" % six.moves.urllib_version pypm_user_agent = get_user_agent(urllib2_user_agent) assert pypm_user_agent req.add_header('User-Agent', pypm_user_agent) try: u = six.moves.urlopen(req) except six.moves.HTTPError as e: if e.code == 304 and use_cache: invoke_start_info('Hit') return False # file not changed in server else: invoke_start_info('Get') raise else: invoke_start_info('Get') bs = 1024*8 # block size; from urllib.py:urlretrieve headers = u.info() # detect total size of the file to be downloaded if 'Content-Length' in headers: total_size = int(headers['Content-Length']) assert total_size >= 0 else: total_size = None total_bytes_transferred = 0 chunk_size = 0 # Hook 1: initialize self.hook_initialize(total_size) with open(target_file + '.part', 'wb') as f: while True: if total_size: assert total_bytes_transferred <= total_size, str(( total_bytes_transferred, total_size)) # Hook 2: transferring self.hook_transferring(chunk_size, total_bytes_transferred) data = u.read(bs) chunk_size = len(data) if len(data) == 0: break f.write(data) total_bytes_transferred += chunk_size sh.mv(target_file + '.part', target_file) # Hook 3: completed self.hook_completed() # save the new headers and other properties if save_properties: custom_dict = props and props.custom or {} # reuse existing props.custom custom_dict['last_attempt_utc'] = BareDateTime.to_string(datetime.utcnow()) urlprops.save(dict(headers), custom_dict) return True # # hooks for various download events # def hook_initialize(self, total_size): """Called just before the download starts - total_size: total size of the URL; None if it is unknown """ def hook_transferring(self, chunk_size, total_bytes_transferred): """Called while the download happens - chunk_size: bytes transferred in this read - total_bytes_transferred: total bytes transferred so far """ def hook_completed(self): """Called when the download is complete and the target file is closed""" class URLProperties(object): """Properties file for the downloaded URL. Named `.{filename}.urlprops`, it is is a JSON dictionary with two keys: - headers: the HTTP headers dictionary - custom: any custom data stored as a dictionary URL property file can be used to retrieve download-related data in future. To automatically create this file, invoke URLGetter like:: >>> URLGetter.get(..., save_properties=True) """ Data = namedtuple('URLProperties_Data', 'headers custom') def __init__(self, target_file): """ - target_file: Local file where the URL is/would be downloaded """ self.target_file = target_file @property def urlprops_filename(self): assert P.isabs(self.target_file), 'not a absolute path: %s' % self.target_file return P.join(P.dirname(self.target_file), '.{basename}.urlprops'.format( basename=P.basename(self.target_file))) def load(self): """Load URL properties, as saved by the `save` method above. Returned value is a namedtuple containing `headers` and `custom` dicts as passed to the `save()` method. Return None, if no target file or urlprops file is found. """ if P.exists(self.target_file): fn = self.urlprops_filename if P.exists(fn): LOG.debug('loading urlprops: %s', fn) with open(fn) as f: data = json.load(f) assert tuple(sorted(data.keys())) == ('custom', 'headers'), \ 'invalid keys: %s' % data.keys() return self.Data(data['headers'], data['custom']) def save(self, headers_dict, custom_dict): """Save headers_dict and custom_dict""" assert P.exists(self.target_file) fn = self.urlprops_filename LOG.debug('writing urlprops: %s', fn) data = dict(headers=headers_dict, custom=custom_dict) json.dumps(data) # catch json errors early (before writing to file) with open(fn, 'w') as f: json.dump(data, f) class Downloader(URLGetter): """An URLGetter that raises ``DownloadError`` on various network errors Useful for catching such errors with one simple try..except.. block. """ def get(self, *args, **kwargs): """Catch network errors and reraise a simple version of them""" try: return super(Downloader, self).get(*args, **kwargs) except six.moves.HTTPError as e: # propagate HTTP events raise except six.moves.URLError as e: raise DownloadError("Failed to download <%s>. Reason: %s" % ( self.url, e)) class DownloaderWithProgress(Downloader): """A Downloader with progress bar""" def hook_initialize(self, total_size): if total_size is None: # XXX: progress bar cannot be used as total_size is unknown # TODO: so, implement a progress bar with optional total_size self.progress = None else: self.progress = textui.ProgressBar( total_size, show_size=_byteshr, note='downloading') def hook_transferring(self, chunk_size, total_bytes_transferred): if self.progress: self.progress.tick(items=chunk_size) def hook_completed(self): if self.progress: self.progress.close() def _byteshr(bytes): """Human-readable version of bytes count""" for x in ['bytes','KB','MB','GB','TB']: if bytes < 1024.0: return "%3.1f%s" % (bytes, x) bytes /= 1024.0 raise ValueError('cannot find human-readable version') def _create_http_basicauth_header(username, password): s = '{0}:{1}'.format(username, password) base64string = encodebytes(s.encode('utf8'))[:-1] authheader = six.b('Basic ') + base64string return authheader