import os
import requests
import six
import sys
import tempfile
import threading
from contextlib
import closing
from six.moves.urllib.parse
import urlparse
from dlmanager
import fs
from dlmanager.persist_limit
import PersistLimit
class DownloadInterrupt(Exception):
"Raised when a download is interrupted."
class Download(object):
"""
Download
is reponsible of downloading one file
in the background.
Example of use: ::
dl = Download(url, dest)
dl.start()
dl.wait()
# this will block until completion / cancel / error
If a download fail
or is canceled, the temporary dest
is removed
from
the disk.
Usually, Downloads are created by using :meth:`DownloadManager.download`.
:param url: the url of the file to download
:param dest: the local file path destination
:param finished_callback: a callback that will be called
in the thread
when the thread work
is done. Takes the download
instance
as a parameter.
:param chunk_size: size of the chunk that will be read. The thread can
not be stopped
while we are reading that chunk size.
:param session: a requests.Session instance that will do do the real
downloading work.
If None, `requests` module
is used.
:param progress: A callable to report the progress (default to
None).
see :meth:`set_progress`.
"""
def __init__(self, url, dest, finished_callback=
None,
chunk_size=16 * 1024, session=
None, progress=
None):
self.thread = threading.Thread(
target=self._download,
args=(url, dest, finished_callback, chunk_size,
session
or requests)
)
self._lock = threading.Lock()
self.__url = url
self.__dest = dest
self.__progress = progress
self.__canceled =
False
self.__error =
None
def start(self):
"""
Start the thread that will do the download.
"""
self.thread.start()
def cancel(self):
"""
Cancel a previously started download.
"""
self.__canceled =
True
def is_canceled(self):
"""
Returns
True if we canceled this download.
"""
return self.__canceled
def is_running(self):
"""
Returns
True if the downloading thread
is running.
"""
return self.thread.is_alive()
def wait(self, raise_if_error=
True):
"""
Block until the downloading thread
is finished.
:param raise_if_error:
if True (the default), :meth:`raise_if_error`
will be called
and raise an error
if any.
"""
while self.thread.is_alive():
try:
# in case of exception here (like KeyboardInterrupt),
# cancel the task.
self.thread.join(0.02)
except:
self.cancel()
raise
# this will raise exception that may happen inside the thread.
if raise_if_error:
self.raise_if_error()
def error(self):
"""
Returns
None or a tuple of three values (type, value, traceback)
that give information about the exception.
"""
return self.__error
def raise_if_error(self):
"""
Raise an error
if any.
If the download was canceled,
raise
:
class:`DownloadInterrupt`.
"""
if self.__error:
six.reraise(*self.__error)
if self.__canceled:
raise DownloadInterrupt()
def set_progress(self, progress):
"""
set a callable to report the progress of the download,
or None to
disable any report.
The callable must take three parameters (download, current, total).
Note that this method
is thread safe, you can call it during a
download.
"""
with self._lock:
self.__progress = progress
def get_dest(self):
"""
Returns the dest.
"""
return self.__dest
def get_url(self):
"""
Returns the url.
"""
return self.__url
def _update_progress(self, current, total):
with self._lock:
if self.__progress:
self.__progress(self, current, total)
def _download(self, url, dest, finished_callback, chunk_size, session):
# save the file under a temporary name
# this allow to not use a broken file in case things went really bad
# while downloading the file (ie the python interpreter is killed
# abruptly)
temp =
None
bytes_so_far = 0
try:
with closing(session.get(url, stream=
True))
as response:
total_size = response.headers.get(
'Content-length',
'').strip()
total_size = int(total_size)
if total_size
else None
self._update_progress(bytes_so_far, total_size)
# we use NamedTemporaryFile as raw open() call was causing
# issues on windows - see:
# https://bugzilla.mozilla.org/show_bug.cgi?id=1185756
with tempfile.NamedTemporaryFile(
delete=
False,
suffix=
'.tmp',
dir=os.path.dirname(dest))
as temp:
for chunk
in response.iter_content(chunk_size):
if self.is_canceled():
break
if chunk:
temp.write(chunk)
bytes_so_far += len(chunk)
self._update_progress(bytes_so_far, total_size)
response.raise_for_status()
except:
self.__error = sys.exc_info()
try:
if temp
is None:
pass # not even opened the temp file, nothing to do
elif self.is_canceled()
or self.__error:
fs.remove(temp.name)
else:
# if all goes well, then rename the file to the real dest
fs.remove(dest)
# just in case it already existed
fs.move(temp.name, dest)
finally:
if finished_callback:
finished_callback(self)
class DownloadManager(object):
"""
DownloadManager
is responsible of starting
and managing downloads inside
a given directory. It will download a file only
if a given filename
is not already there.
Note that background downloads needs to be stopped.
For example,
if
you have an exception
while a download
is occuring, python will only
exit when the download will finish. To get rid of that, there
is a
possible idiom: ::
def download_things(manager):
# do things with the manager
manager.download(url1, f1)
manager.download(url2, f2)
...
manager = DownloadManager(destdir)
try:
download_things(manager)
finally:
# ensure we cancel all background downloads to ask the end
# of possible remainings threads
manager.cancel()
:param destdir: a directory where files are downloaded. It will be created
if it does
not exists.
:param session: a requests session.
If None, one will be created
for you.
:param persist_limit: an instance of :
class:`PersistLimit`, to allow
limiting the size of the download dir. Defaults
to
None, meaning no limit.
"""
def __init__(self, destdir, session=
None, persist_limit=
None):
self.destdir = destdir
self.session = session
or requests.Session()
self._downloads = {}
self._lock = threading.Lock()
self.persist_limit = persist_limit
or PersistLimit(0)
self.persist_limit.register_dir_content(self.destdir)
# if persist folder does not exist, create it
if not os.path.isdir(destdir):
os.makedirs(destdir)
def get_dest(self, fname):
return os.path.join(self.destdir, fname)
def cancel(self, cancel_if=
None):
"""
Cancel downloads,
if any.
if cancel_if
is given, it must be a callable that take the download
instance
as parameter,
and return True if the download needs to be
canceled.
Note that download threads won
't be stopped directly.
"""
with self._lock:
for download
in six.itervalues(self._downloads):
if cancel_if
is None or cancel_if(download):
if download.is_running():
download.cancel()
def wait(self, raise_if_error=
True):
"""
Wait
for all downloads to be finished.
"""
for download
in self._downloads.values():
download.wait(raise_if_error=raise_if_error)
def download(self, url, fname=
None, progress=
None):
"""
Returns a started :
class:`Download` instance,
or None if fname
is
already present
in destdir.
if a download
is already running
for the given fname, it
is just
returned.
Else the download
is created, started
and returned.
:param url: url of the file to download.
:param fname: name to give
for the downloaded file.
If None, it will
be the name extracted
in the url.
:param progress: a callable to report the download progress,
or None.
See :meth:`Download.set_progress`.
"""
if fname
is None:
fname = urlparse(url).path.split(
'/')[-1]
dest = self.get_dest(fname)
with self._lock:
# if we are downloading, returns the instance
if dest
in self._downloads:
dl = self._downloads[dest]
if progress:
dl.set_progress(progress)
return dl
if os.path.exists(dest):
return None
# else create the download (will be automatically removed of
# the list on completion) start it, and returns that.
with self._lock:
download = Download(url, dest,
session=self.session,
finished_callback=self._download_finished,
progress=progress)
self._downloads[dest] = download
download.start()
self._download_started(download)
return download
def _download_started(self, dl):
"""
Useful when sub-classing. Report the start event of a download.
:param dl: The :
class:`Download` instance.
"""
pass
def _download_finished(self, dl):
"""
Useful when sub-classing. Report the end of a download.
Note that this
is executed
in the download thread. Also, you should
make sure to call the base implementation.
:param dl: The :
class:`Download` instance.
"""
with self._lock:
dest = dl.get_dest()
del self._downloads[dest]
self.persist_limit.register_file(dest)
self.persist_limit.remove_old_files()
