diff --git a/youtube_dlc/downloader/common.py b/youtube_dlc/downloader/common.py index 7d303be1c..05d1dd5f7 100644 --- a/youtube_dlc/downloader/common.py +++ b/youtube_dlc/downloader/common.py @@ -7,6 +7,7 @@ import time import random from ..compat import compat_os_name +from ..heartbeat import Heartbeat from ..utils import ( decodeArgument, encodeFilename, @@ -372,6 +373,12 @@ class FileDownloader(object): '[download] Sleeping %s seconds...' % ( sleep_interval_sub)) time.sleep(sleep_interval_sub) + + if info_dict.get('heartbeat'): + self.heartbeat = Heartbeat(self.ydl, info_dict.get('heartbeat')) + self.add_progress_hook(self.heartbeat.check_download_status) + self.heartbeat.start() + return self.real_download(filename, info_dict) def real_download(self, filename, info_dict): diff --git a/youtube_dlc/extractor/niconico.py b/youtube_dlc/extractor/niconico.py index eb07ca776..1c6172794 100644 --- a/youtube_dlc/extractor/niconico.py +++ b/youtube_dlc/extractor/niconico.py @@ -254,6 +254,18 @@ class NiconicoIE(InfoExtractor): } }).encode()) + heartbeat_url = '{}/{}?_format=json&_method=PUT'.format(session_api_endpoint['url'], session_response['data']['session']['id']) + heartbeat_headers = { + 'Accept': 'application/json', + 'Content-Type': 'application/json', + } + heartbeat = { + 'data': json.dumps(session_response['data']), + 'headers': heartbeat_headers, + 'interval': session_api_data['heartbeat_lifetime'] / 2000, + 'url': heartbeat_url, + } + resolution = video_quality.get('resolution', {}) return { @@ -264,6 +276,7 @@ class NiconicoIE(InfoExtractor): 'vbr': float_or_none(video_quality.get('bitrate'), 1000), 'height': resolution.get('height'), 'width': resolution.get('width'), + 'heartbeat': heartbeat, } def _real_extract(self, url): diff --git a/youtube_dlc/heartbeat.py b/youtube_dlc/heartbeat.py new file mode 100644 index 000000000..0b0bb838c --- /dev/null +++ b/youtube_dlc/heartbeat.py @@ -0,0 +1,58 @@ +# coding: utf-8 +from __future__ import unicode_literals + +import threading +import traceback + +from .utils import ( + compat_str, + encode_compat_str, + sanitized_Request +) + + +class Heartbeat(object): + def __init__(self, ydl, params): + self.ydl = ydl + + data = params.get('data') + if isinstance(data, compat_str): + data = data.encode() + # Python 2 does not allow us to set HTTP method + # it is POST if Request has data, otherwise GET + self.request = sanitized_Request( + params.get('url'), + data=data, + headers=params.get('headers', {}) + ) + + self.interval = params.get('interval', 30) + self.cancelled = False + self.parent_thread = threading.current_thread() + self.thread = threading.Thread(target=self.__heartbeat) + + def start(self): + self.ydl.to_screen('[heartbeat] Heartbeat every %s seconds' % self.interval) + self.thread.start() + + def cancel(self): + self.cancelled = True + + def check_download_status(self, progress): + status = progress.get('status') + if status == 'finished' or status == 'error': + self.cancel() + + def __heartbeat(self): + while not self.cancelled: + try: + if self.ydl.params.get('verbose'): + self.ydl.to_screen('[heartbeat]') + self.ydl.urlopen(self.request) + except Exception: + self.ydl.report_warning("Heartbeat failed") + if self.ydl.params.get('verbose'): + self.ydl.to_stderr(encode_compat_str(traceback.format_exc())) + self.parent_thread.join(self.interval) + if not self.parent_thread.is_alive(): + break