Adding dowload limit feature

This commit is contained in:
Shadow035 2020-04-20 14:57:21 +05:30
parent 00eb865b3c
commit 640a3a28b9
6 changed files with 61 additions and 0 deletions

3
.gitignore vendored
View File

@ -51,3 +51,6 @@ venv/
# VS Code related files
.vscode
# Visual Studio related files
.vs/

View File

@ -65,6 +65,7 @@ from .utils import (
locked_file,
make_HTTPS_handler,
MaxDownloadsReached,
MaxDataReached,
orderedSet,
PagedList,
parse_filesize,
@ -340,6 +341,9 @@ class YoutubeDL(object):
_num_downloads = None
_screen_file = None
global _downloaded_data
_downloaded_data = None
def __init__(self, params=None, auto_init=True):
"""Create a FileDownloader object with the given options."""
if params is None:
@ -350,6 +354,7 @@ class YoutubeDL(object):
self._progress_hooks = []
self._download_retcode = 0
self._num_downloads = 0
self._downloaded_data = 0
self._screen_file = [sys.stdout, sys.stderr][params.get('logtostderr', False)]
self._err_file = sys.stderr
self.params = {
@ -821,6 +826,8 @@ class YoutubeDL(object):
break
except MaxDownloadsReached:
raise
except MaxDataReached:
raise
except Exception as e:
if self.params.get('ignoreerrors', False):
self.report_error(error_to_compat_str(e), tb=encode_compat_str(traceback.format_exc()))
@ -2022,6 +2029,9 @@ class YoutubeDL(object):
except MaxDownloadsReached:
self.to_screen('[info] Maximum number of downloaded files reached.')
raise
except MaxDataReached:
self.to_screen('[info] Maximum data limit reached.')
raise
else:
if self.params.get('dump_single_json', False):
self.to_stdout(json.dumps(res))

View File

@ -28,6 +28,7 @@ from .utils import (
expand_path,
match_filter_func,
MaxDownloadsReached,
MaxDataReached,
preferredencoding,
read_batch_urls,
SameFileError,
@ -150,6 +151,16 @@ def _real_main(argv=None):
if numeric_limit is None:
parser.error('invalid rate limit specified')
opts.ratelimit = numeric_limit
if opts.max_data is not None:
numeric_limit = FileDownloader.parse_bytes(opts.max_data)
if numeric_limit is None:
parser.error('invalid max_data specified')
opts.max_data = numeric_limit
if opts.max_data_new is not None:
numeric_limit = FileDownloader.parse_bytes(opts.max_data_new)
if numeric_limit is None:
parser.error('invalid max_data_new specified')
opts.max_data_new = numeric_limit
if opts.min_filesize is not None:
numeric_limit = FileDownloader.parse_bytes(opts.min_filesize)
if numeric_limit is None:
@ -386,6 +397,8 @@ def _real_main(argv=None):
'write_pages': opts.write_pages,
'test': opts.test,
'keepvideo': opts.keepvideo,
'max_data': opts.max_data,
'max_data_new': opts.max_data_new,
'min_filesize': opts.min_filesize,
'max_filesize': opts.max_filesize,
'min_views': opts.min_views,
@ -465,6 +478,9 @@ def _real_main(argv=None):
except MaxDownloadsReached:
ydl.to_screen('--max-download limit reached, aborting.')
retcode = 101
except MaxDataReached:
ydl.to_screen('--max-data limit reached, aborting.')
retcode = 102
sys.exit(retcode)

View File

@ -21,6 +21,7 @@ from ..utils import (
write_xattr,
XAttrMetadataError,
XAttrUnavailableError,
MaxDataReached
)
@ -187,6 +188,8 @@ class HttpFD(FileDownloader):
raise RetryDownload(err)
def download():
global _downloaded_data
data_len = ctx.data.info().get('Content-length', None)
# Range HTTP header may be ignored/unsupported by a webserver
@ -197,6 +200,9 @@ class HttpFD(FileDownloader):
if is_test and (data_len is None or int(data_len) > self._TEST_FILE_SIZE):
data_len = self._TEST_FILE_SIZE
max_datalimit = self.params.get('max_data')
max_datalimit_new = self.params.get('max_data_new')
if data_len is not None:
data_len = int(data_len) + ctx.resume_len
min_data_len = self.params.get('min_filesize')
@ -207,6 +213,9 @@ class HttpFD(FileDownloader):
if max_data_len is not None and data_len > max_data_len:
self.to_screen('\r[download] File is larger than max-filesize (%s bytes > %s bytes). Aborting.' % (data_len, max_data_len))
return False
if max_datalimit_new is not None and _downloaded_data + data_len > max_datalimit_new:
self.to_screen('\r[download] Next file size (%s bytes), Maxmimum data limit (%s bytes) expected. Aborting.' % (data_len, max_datalimit_new))
raise MaxDataReached()
byte_counter = 0 + ctx.resume_len
block_size = ctx.block_size
@ -299,6 +308,16 @@ class HttpFD(FileDownloader):
'elapsed': now - ctx.start_time,
})
_downloaded_data += len(data_block)
if max_datalimit is not None:
if _downloaded_data > max_datalimit:
self.to_screen('\r[download] Maxmimum data limit (%s bytes) reached. Aborting.' % (data_len, max_datalimit))
raise MaxDataReached()
elif max_datalimit_new is not None and _downloaded_data > max_datalimit_new:
self.to_screen('\r[download] Maxmimum data limit (%s bytes) reached. Aborting.' % (data_len, max_datalimit_new))
raise MaxDataReached()
if is_test and byte_counter == data_len:
break

View File

@ -279,6 +279,14 @@ def parseOpts(overrideArguments=None):
'--max-downloads',
dest='max_downloads', metavar='NUMBER', type=int, default=None,
help='Abort after downloading NUMBER files')
selection.add_option(
'--max-data',
metavar='SIZE', dest='max_data', default=None,
help='Stop download if total downloaded data reached SIZE (e.g. 50k or 44.6m)')
selection.add_option(
'--max-data-new',
metavar='SIZE', dest='max_data_new', default=None,
help='Stop downloading new file if estimated total downloaded data exceeds SIZE (e.g. 50k or 44.6m)')
selection.add_option(
'--min-filesize',
metavar='SIZE', dest='min_filesize', default=None,

View File

@ -2416,6 +2416,11 @@ class MaxDownloadsReached(YoutubeDLError):
pass
class MaxDataReached(YoutubeDLError):
""" --max-data limit has been reached. """
pass
class UnavailableVideoError(YoutubeDLError):
"""Unavailable Format exception.