From 21d93613a231670c58176b80aae23c95beddb9d5 Mon Sep 17 00:00:00 2001
From: taizan-hokuto <55448286+taizan-hokuto@users.noreply.github.com>
Date: Sat, 14 Mar 2020 08:00:31 +0900
Subject: [PATCH] Handling JSONDecodeError

---
 pytchat/exceptions.py           | 22 +++++++++++++-------
 pytchat/tool/extract/asyncdl.py | 37 ++++++++++++++++++++++++++-------
 2 files changed, 44 insertions(+), 15 deletions(-)

diff --git a/pytchat/exceptions.py b/pytchat/exceptions.py
index cffb955..11e5ba1 100644
--- a/pytchat/exceptions.py
+++ b/pytchat/exceptions.py
@@ -1,46 +1,52 @@
 class ChatParseException(Exception):
     '''
-    チャットデータをパースするライブラリが投げる例外の基底クラス
+    Base exception thrown by the parser
     '''
     pass
 
 class NoYtinitialdataException(ChatParseException):
     '''
-    配信ページ内にチャットデータurlが見つからないときに投げる例外
+    Thrown when the video is not found.
     '''
     pass
 
 class ResponseContextError(ChatParseException):
     '''
-    配信ページでチャットデータ無効の時に投げる例外
+    Thrown when chat data is invalid.
     '''
     pass
 
 class NoLivechatRendererException(ChatParseException):
     '''
-    チャットデータのJSON中にlivechatRendererがない時に投げる例外
+    Thrown when livechatRenderer is missing in JSON.
     '''
     pass
 
 
 class NoContentsException(ChatParseException):
     '''
-    チャットデータのJSON中にContinuationContentsがない時に投げる例外
+    Thrown when ContinuationContents is missing in JSON.
     '''
     pass
 
 class NoContinuationsException(ChatParseException):
     '''
-    チャットデータのContinuationContents中にcontinuationがない時に投げる例外
+    Thrown when continuation is missing in ContinuationContents.
     '''
     pass
 
 class IllegalFunctionCall(Exception):
     '''
-    set_callback()を実行済みにもかかわらず
-    get()を呼び出した場合の例外
+    Thrown when get () is called even though 
+    set_callback () has been executed.
     '''
     pass
 
 class InvalidVideoIdException(Exception):
+    '''
+    Thrown when the video_id is not exist (VideoInfo).
+    '''
     pass
+
+class UnknownConnectionError(Exception):
+    pass
\ No newline at end of file
diff --git a/pytchat/tool/extract/asyncdl.py b/pytchat/tool/extract/asyncdl.py
index c361754..084f037 100644
--- a/pytchat/tool/extract/asyncdl.py
+++ b/pytchat/tool/extract/asyncdl.py
@@ -7,12 +7,15 @@ from . worker import ExtractWorker
 from . patch import Patch
 from ... import config 
 from ... paramgen import arcparam
+from ... exceptions import UnknownConnectionError
 from concurrent.futures import CancelledError
+from json import JSONDecodeError
 from urllib.parse import quote
 
 headers = config.headers
 REPLAY_URL = "https://www.youtube.com/live_chat_replay/" \
              "get_live_chat_replay?continuation="
+MAX_RETRY_COUNT = 3
 
 def _split(start, end, count, min_interval_sec = 120):
     """
@@ -53,13 +56,22 @@ def ready_blocks(video_id, duration, div, callback):
             tasks = [_create_block(session, video_id,  seektime, callback)
                 for  seektime in _split(-1, duration, div)]
             return await asyncio.gather(*tasks)
-
+    
     async def _create_block(session, video_id, seektime, callback):
         continuation = arcparam.getparam(video_id, seektime = seektime)
         url = f"{REPLAY_URL}{quote(continuation)}&pbj=1"
-        async with session.get(url, headers = headers) as resp:
-            text = await resp.text()
-        next_continuation, actions = parser.parse(json.loads(text))
+        for _ in range(MAX_RETRY_COUNT):
+            try :
+                async with session.get(url, headers = headers) as resp:
+                    text = await resp.text()
+                next_continuation, actions = parser.parse(json.loads(text))
+                break
+            except JSONDecodeError:
+                await asyncio.sleep(3)
+        else:
+            cancel()
+            raise UnknownConnectionError("Abort: Unknown connection error.")
+
         if actions:
             first = parser.get_offset(actions[0])
             last = parser.get_offset(actions[-1])
@@ -71,6 +83,7 @@ def ready_blocks(video_id, duration, div, callback):
                 first = first,
                 last = last
             )
+            
     """
     fetch initial blocks.
     """  
@@ -95,9 +108,18 @@ def fetch_patch(callback, blocks, video_id):
 
     async def _fetch(continuation,session) -> Patch:
         url = f"{REPLAY_URL}{quote(continuation)}&pbj=1"
-        async with session.get(url,headers = config.headers) as resp:
-            chat_json = await resp.text()
-        continuation, actions = parser.parse(json.loads(chat_json))
+        for _ in range(MAX_RETRY_COUNT):
+            try:
+                async with session.get(url,headers = config.headers) as resp:
+                    chat_json = await resp.text()
+                continuation, actions = parser.parse(json.loads(chat_json))
+                break
+            except JSONDecodeError:
+                await asyncio.sleep(3)
+        else:
+            cancel()
+            raise UnknownConnectionError("Abort: Unknown connection error.")
+
         if actions:
             last = parser.get_offset(actions[-1])
             first = parser.get_offset(actions[0])
@@ -105,6 +127,7 @@ def fetch_patch(callback, blocks, video_id):
                 callback(actions, last - first)
             return Patch(actions, continuation, first, last)
         return Patch(continuation = continuation)
+
     """
     allocate workers and assign blocks.
     """