yt_dlp/downloader/youtube_live_chat.py

   1 from __future__ import division, unicode_literals
   2
   3 import json
   4
   5 from .fragment import FragmentFD
   6 from ..compat import compat_urllib_error
   7 from ..utils import (
   8     try_get,
   9     RegexNotFoundError,
  10 )
  11 from ..extractor.youtube import YoutubeBaseInfoExtractor as YT_BaseIE
  12
  13
  14 class YoutubeLiveChatReplayFD(FragmentFD):
  15     """ Downloads YouTube live chat replays fragment by fragment """
  16
  17     FD_NAME = 'youtube_live_chat_replay'
  18
  19     def real_download(self, filename, info_dict):
  20         video_id = info_dict['video_id']
  21         self.to_screen('[%s] Downloading live chat' % self.FD_NAME)
  22
  23         fragment_retries = self.params.get('fragment_retries', 0)
  24         test = self.params.get('test', False)
  25
  26         ctx = {
  27             'filename': filename,
  28             'live': True,
  29             'total_frags': None,
  30         }
  31
  32         ie = YT_BaseIE(self.ydl)
  33
  34         def dl_fragment(url, data=None, headers=None):
  35             http_headers = info_dict.get('http_headers', {})
  36             if headers:
  37                 http_headers = http_headers.copy()
  38                 http_headers.update(headers)
  39             return self._download_fragment(ctx, url, info_dict, http_headers, data)
  40
  41         def download_and_parse_fragment(url, frag_index, request_data):
  42             count = 0
  43             while count <= fragment_retries:
  44                 try:
  45                     success, raw_fragment = dl_fragment(url, request_data, {'content-type': 'application/json'})
  46                     if not success:
  47                         return False, None, None
  48                     try:
  49                         data = ie._extract_yt_initial_data(video_id, raw_fragment.decode('utf-8', 'replace'))
  50                     except RegexNotFoundError:
  51                         data = None
  52                     if not data:
  53                         data = json.loads(raw_fragment)
  54                     live_chat_continuation = try_get(
  55                         data,
  56                         lambda x: x['continuationContents']['liveChatContinuation'], dict) or {}
  57                     offset = continuation_id = None
  58                     processed_fragment = bytearray()
  59                     for action in live_chat_continuation.get('actions', []):
  60                         if 'replayChatItemAction' in action:
  61                             replay_chat_item_action = action['replayChatItemAction']
  62                             offset = int(replay_chat_item_action['videoOffsetTimeMsec'])
  63                         processed_fragment.extend(
  64                             json.dumps(action, ensure_ascii=False).encode('utf-8') + b'\n')
  65                     if offset is not None:
  66                         continuation_id = try_get(
  67                             live_chat_continuation,
  68                             lambda x: x['continuations'][0]['liveChatReplayContinuationData']['continuation'])
  69                     self._append_fragment(ctx, processed_fragment)
  70
  71                     return True, continuation_id, offset
  72                 except compat_urllib_error.HTTPError as err:
  73                     count += 1
  74                     if count <= fragment_retries:
  75                         self.report_retry_fragment(err, frag_index, count, fragment_retries)
  76             if count > fragment_retries:
  77                 self.report_error('giving up after %s fragment retries' % fragment_retries)
  78                 return False, None, None
  79
  80         self._prepare_and_start_frag_download(ctx)
  81
  82         success, raw_fragment = dl_fragment(info_dict['url'])
  83         if not success:
  84             return False
  85         try:
  86             data = ie._extract_yt_initial_data(video_id, raw_fragment.decode('utf-8', 'replace'))
  87         except RegexNotFoundError:
  88             return False
  89         continuation_id = try_get(
  90             data,
  91             lambda x: x['contents']['twoColumnWatchNextResults']['conversationBar']['liveChatRenderer']['continuations'][0]['reloadContinuationData']['continuation'])
  92         # no data yet but required to call _append_fragment
  93         self._append_fragment(ctx, b'')
  94
  95         ytcfg = ie._extract_ytcfg(video_id, raw_fragment.decode('utf-8', 'replace'))
  96
  97         if not ytcfg:
  98             return False
  99         api_key = try_get(ytcfg, lambda x: x['INNERTUBE_API_KEY'])
 100         innertube_context = try_get(ytcfg, lambda x: x['INNERTUBE_CONTEXT'])
 101         if not api_key or not innertube_context:
 102             return False
 103         url = 'https://www.youtube.com/youtubei/v1/live_chat/get_live_chat_replay?key=' + api_key
 104
 105         frag_index = offset = 0
 106         while continuation_id is not None:
 107             frag_index += 1
 108             request_data = {
 109                 'context': innertube_context,
 110                 'continuation': continuation_id,
 111             }
 112             if frag_index > 1:
 113                 request_data['currentPlayerState'] = {'playerOffsetMs': str(max(offset - 5000, 0))}
 114             success, continuation_id, offset = download_and_parse_fragment(
 115                 url, frag_index, json.dumps(request_data, ensure_ascii=False).encode('utf-8') + b'\n')
 116             if not success:
 117                 return False
 118             if test:
 119                 break
 120
 121         self._finish_frag_download(ctx)
 122         return True