Update download method

====  Tech Notes ====
- Fix cleanup
- Add support for custom header
- create directory as needed
pull/1150/head^2
Åke Forslund 2017-10-16 16:08:03 +02:00 committed by Steve Penrod
parent 7990e168c8
commit 3deb65b0c2
1 changed files with 21 additions and 9 deletions

View File

@ -16,7 +16,7 @@ from threading import Thread
import os
import requests
from os.path import exists
from os.path import exists, dirname
import subprocess
_running_downloads = {}
@ -52,7 +52,7 @@ class Downloader(Thread):
`func(dest)`
"""
def __init__(self, url, dest, complete_action=None):
def __init__(self, url, dest, complete_action=None, header=None):
super(Downloader, self).__init__()
self.url = url
self.dest = dest
@ -60,21 +60,32 @@ class Downloader(Thread):
self.status = None
self.done = False
self._abort = False
self.header = header
# Create directories as needed
if not exists(dirname(dest)):
os.makedirs(dirname(dest))
# Start thread
self.daemon = True
self.start()
def perform_download(self, dest):
cmd = ['wget', '-c', self.url, '-O', dest,
'--tries=20', '--read-timeout=5']
if self.header:
cmd += ['--header={}'.format(self.header)]
return subprocess.call(cmd)
def run(self):
"""
Does the actual download.
"""
r = requests.get(self.url, stream=True)
tmp = _get_download_tmp(self.dest)
self.status = self.perform_download(tmp)
self.status = subprocess.call(['wget', '-c', self.url, '-O', tmp,
'--tries=20', '--read-timeout=5'])
print self.status
if not self._abort and self.status == 0:
self.finalize(tmp)
else:
@ -95,7 +106,7 @@ class Downloader(Thread):
if self.complete_action:
self.complete_action(self.dest)
def cleanup(tmp):
def cleanup(self, tmp):
"""
Cleanup after download attempt
"""
@ -111,9 +122,10 @@ class Downloader(Thread):
self._abort = True
def download(url, dest, complete_action=None):
def download(url, dest, complete_action=None, header=None):
global _running_downloads
arg_hash = hash(url + dest)
if arg_hash not in _running_downloads:
_running_downloads[arg_hash] = Downloader(url, dest, complete_action)
_running_downloads[arg_hash] = Downloader(url, dest, complete_action,
header)
return _running_downloads[arg_hash]