fmt/support/download.py

# A file downloader.

import contextlib, os, tempfile, timer, urllib2, urlparse

class Downloader:
    def __init__(self, dir=None):
        self.dir = dir

    # Downloads a file and removes it when exiting a block.
    # Usage:
    #  d = Downloader()
    #  with d.download(url) as f:
    #    use_file(f)
    def download(self, url, cookie=None):
        suffix = os.path.splitext(urlparse.urlsplit(url)[2])[1]
        fd, filename = tempfile.mkstemp(suffix=suffix, dir=self.dir)
        os.close(fd)
        with timer.print_time('Downloading', url, 'to', filename):
            opener = urllib2.build_opener()
            if cookie:
                opener.addheaders.append(('Cookie', cookie))
            num_tries = 2
            for i in range(num_tries):
                try:
                    f = opener.open(url)
                except urllib2.URLError, e:
                    print('Failed to open url', url)
                    continue
                length = f.headers.get('content-length')
                if not length:
                    print('Failed to get content-length')
                    continue
                length = int(length)
                with open(filename, 'wb') as out:
                    count = 0
                    while count < length:
                        data = f.read(1024 * 1024)
                        count += len(data)
                        out.write(data)
        @contextlib.contextmanager
        def remove(filename):
            try:
                yield filename
            finally:
                os.remove(filename)
        return remove(filename)