This commit is contained in:
insaneracist 2021-08-24 16:33:24 -04:00 committed by GitHub
commit 16760b4412
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 78 additions and 0 deletions

View File

@ -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):

View File

@ -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):

58
youtube_dlc/heartbeat.py Normal file
View File

@ -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