From 6fdb3bf8cf3d4f0475caa45354ae68d8638d504b Mon Sep 17 00:00:00 2001
From: taizan-hokuto <55448286+taizan-hokuto@users.noreply.github.com>
Date: Thu, 13 Feb 2020 21:15:39 +0900
Subject: [PATCH] Implement cancell download

---
 pytchat/tool/asyncdl.py    | 23 ++++++++++++++++++++++-
 pytchat/tool/downloader.py |  3 +++
 2 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/pytchat/tool/asyncdl.py b/pytchat/tool/asyncdl.py
index 708f74b..6866c8a 100644
--- a/pytchat/tool/asyncdl.py
+++ b/pytchat/tool/asyncdl.py
@@ -8,6 +8,7 @@ from . dlworker import DownloadWorker
 from .. paramgen import arcparam
 from .. import config 
 from urllib.parse import quote
+from concurrent.futures import CancelledError
 
 headers = config.headers
 REPLAY_URL = "https://www.youtube.com/live_chat_replay/" \
@@ -127,4 +128,24 @@ def download_chunk(callback, blocks, video_id):
         return [], continuation, None, None
     
     loop = asyncio.get_event_loop()
-    loop.run_until_complete(_allocate_workers())
+    try:
+        loop.run_until_complete(_allocate_workers())
+    except CancelledError:
+        pass
+
+
+async def shutdown():
+    print("\nshutdown...")
+    tasks = [t for t in asyncio.all_tasks()
+         if t is not asyncio.current_task()]
+    for task in tasks:
+        task.cancel()
+        try:
+            await task
+        except asyncio.CancelledError:
+            pass
+
+def cancel():
+    loop = asyncio.get_event_loop()
+    loop.create_task(shutdown())
+    
\ No newline at end of file
diff --git a/pytchat/tool/downloader.py b/pytchat/tool/downloader.py
index 07a4178..2543c44 100644
--- a/pytchat/tool/downloader.py
+++ b/pytchat/tool/downloader.py
@@ -84,3 +84,6 @@ def download(video_id, div = 1, callback = None, processor = None):
         [{'video_id':None,'timeout':1,'chatdata' : [action
         ["replayChatItemAction"]["actions"][0] for action in data]}]
     )
+
+def cancel():
+    asyncdl.cancel()