Merge pull request #266 from Guovin/dev

Dev

main.py

@@ -12,6 +12,7 @@ from utils.tools import (
     get_ip_address,
     convert_to_m3u,
     get_result_file_content,
+    merge_objects,
 )
 from updates.subscribe import get_channels_by_subscribe_urls
 from updates.multicast import get_channels_by_multicast
@@ -145,8 +146,9 @@ class UpdateSource:
                 )
                 self.start_time = time()
                 self.pbar = tqdm_asyncio(total=self.total, desc="Sorting")
+                self.sort_n = 0
                 self.channel_data = await process_sort_channel_list(
-                    self.channel_data, self.sort_pbar_update
+                    data=self.channel_data, callback=self.sort_pbar_update
                 )
             no_result_cate_names = [
                 (cate, name)
@@ -170,23 +172,28 @@ class UpdateSource:
                 self.total = len(
                     [name for obj in sup_channel_items.values() for name in obj.keys()]
                 )
-                if config.getboolean("Settings", "open_sort"):
+                if self.total > 0 and config.getboolean("Settings", "open_sort"):
                     self.update_progress(
                         f"正在对补充频道测速排序, 共{self.total}个频道",
                         0,
                     )
                     self.start_time = time()
                     self.pbar = tqdm_asyncio(total=self.total, desc="Sorting")
-                    self.channel_data = await process_sort_channel_list(
-                        self.channel_data,
-                        self.sort_pbar_update,
+                    self.sort_n = 0
+                    sup_channel_items = await process_sort_channel_list(
+                        data=sup_channel_items,
+                        callback=self.sort_pbar_update,
                     )
+                    self.channel_data = merge_objects(
+                        self.channel_data, sup_channel_items
+                    )
+            self.total = len(channel_names)
             self.pbar = tqdm(total=self.total, desc="Writing")
             self.start_time = time()
             write_channel_to_file(
-                channel_items_obj_items,
-                self.channel_data,
-                lambda: self.pbar_update(name="写入结果"),
+                items=channel_items_obj_items,
+                data=self.channel_data,
+                callback=lambda: self.pbar_update(name="写入结果"),
             )
             self.pbar.close()
             user_final_file = config.get("Settings", "final_file")

utils/channel.py

@@ -591,7 +591,9 @@ def append_all_method_data_keep_all(
     return data
 
 
-async def sort_channel_list(semaphore, cate, name, info_list, is_ffmpeg, callback):
+async def sort_channel_list(
+    semaphore=None, cate=None, name=None, info_list=None, ffmpeg=False, callback=None
+):
     """
     Sort the channel list
     """
@@ -600,7 +602,7 @@ async def sort_channel_list(semaphore, cate, name, info_list, is_ffmpeg, callbac
         try:
             if info_list:
                 sorted_data = await sort_urls_by_speed_and_resolution(
-                    info_list, is_ffmpeg
+                    data=info_list, ffmpeg=ffmpeg
                 )
                 if sorted_data:
                     for (
@@ -618,11 +620,12 @@ async def sort_channel_list(semaphore, cate, name, info_list, is_ffmpeg, callbac
         except Exception as e:
             logging.error(f"Error: {e}")
         finally:
-            callback()
+            if callback:
+                callback()
             return {"cate": cate, "name": name, "data": data}
 
 
-async def process_sort_channel_list(channel_data, callback):
+async def process_sort_channel_list(data=None, callback=None):
     """
     Processs the sort channel list
     """
@@ -635,31 +638,29 @@ async def process_sort_channel_list(channel_data, callback):
     tasks = [
         asyncio.create_task(
             sort_channel_list(
-                semaphore,
-                cate,
-                name,
-                info_list,
-                is_ffmpeg,
-                lambda: callback(),
+                semaphore=semaphore,
+                cate=cate,
+                name=name,
+                info_list=info_list,
+                ffmpeg=is_ffmpeg,
+                callback=callback,
             )
         )
-        for cate, channel_obj in channel_data.items()
+        for cate, channel_obj in data.items()
         for name, info_list in channel_obj.items()
     ]
     sort_results = await tqdm_asyncio.gather(*tasks, desc="Sorting")
-    channel_data = {}
+    data = {}
     for result in sort_results:
         if result:
             cate = result.get("cate")
             name = result.get("name")
-            data = result.get("data")
-            channel_data = append_data_to_info_data(
-                channel_data, cate, name, data, False
-            )
-    return channel_data
+            result_data = result.get("data")
+            data = append_data_to_info_data(data, cate, name, result_data, False)
+    return data
 
 
-def write_channel_to_file(items, data, callback):
+def write_channel_to_file(items=None, data=None, callback=None):
     """
     Write channel to file
     """
@@ -671,7 +672,8 @@ def write_channel_to_file(items, data, callback):
                 print("write:", cate, name, "num:", len(channel_urls))
                 update_channel_urls_txt(cate, name, channel_urls)
             finally:
-                callback()
+                if callback:
+                    callback()
     for handler in logging.root.handlers[:]:
         handler.close()
         logging.root.removeHandler(handler)

utils/speed.py

@@ -129,23 +129,19 @@ async def get_info_with_speed(url_info):
         return float("inf")
 
 
-async def sort_urls_by_speed_and_resolution(infoList, ffmpeg=False):
+async def sort_urls_by_speed_and_resolution(data=None, ffmpeg=False):
     """
     Sort by speed and resolution
     """
     if ffmpeg:
         response = await asyncio.gather(
-            *(get_info_with_speed(url_info) for url_info in infoList)
+            *(get_info_with_speed(url_info) for url_info in data)
         )
         valid_response = [res for res in response if res != float("inf")]
     else:
-        response_times = await asyncio.gather(
-            *(get_speed(url) for url, _, _ in infoList)
-        )
+        response_times = await asyncio.gather(*(get_speed(url) for url, _, _ in data))
         valid_response = [
-            (info, rt)
-            for info, rt in zip(infoList, response_times)
-            if rt != float("inf")
+            (info, rt) for info, rt in zip(data, response_times) if rt != float("inf")
         ]
 
     def extract_resolution(resolution_str):