move redis to separate module

README.md

@@ -15,4 +15,4 @@ Python获取bilibili直播弹幕的库，使用WebSocket协议，支持web端和
     pip install -r requirements.txt
     ```
 
-3. 例程看[sample.py](./sample.py)和[open_live_sample.py](./open_live_sample.py)
+3. 例程看[sample.py](./main.py)和[open_live_sample.py](./open_live_sample.py)

main.py

@@ -9,43 +9,17 @@ import time
 from typing import *
 
 import aiohttp
-import redis
+import redis_util as redis
 
 import blivedm
 from blivedm.clients import ws_base
 
-try:
-    with open("config.json") as conf_f:
-        conf = json.load(conf_f)
-except FileNotFoundError:
-    conf = {}
-redis_conf = conf.get("redis", {})
 
-os.makedirs("logs", exist_ok=True)
-room_status_log = open("logs/room_status.jsonl", "a", encoding="utf-8")
-
-# set log level to INFO
-logging.basicConfig(level=logging.INFO)
-
-logger = logging.getLogger()
-
-room_list_key = "bilibili.live.danmu.room_list"
-r = redis.Redis(
-    host=redis_conf.get("host", "127.0.0.1"),
-    port=redis_conf.get("port", 6379),
-    db=redis_conf.get("port", 0),
-)
-
-TEST_ROOM_IDS = [int(room_id) for room_id in set(r.lrange(room_list_key, 0, -1))]
-clients = {}
-log_files = {}
-
-with open("cookie.txt", "r", encoding="utf-8") as f:
-    SESSDATA = f.read().strip()
-
-session: Optional[aiohttp.ClientSession] = None
+def room_from_config():
+    return conf.get("room_ids", [])
 
 
+# noinspection PyUnusedLocal
 def handle_sighup(signum, frame):
     global log_files
     global room_status_log
@@ -64,12 +38,37 @@ def handle_sighup(signum, frame):
     room_status_log = open("logs/room_status.jsonl", "a", encoding="utf-8")
 
 
-signal.signal(signal.SIGHUP, handle_sighup)
+def on_room_changed(new_room_ids):
+    global TEST_ROOM_IDS
+
+    TEST_ROOM_IDS = list(set(new_room_ids + room_from_config()))
+
+    for removed in clients.keys() - TEST_ROOM_IDS:
+        asyncio.run(clients[removed].stop_and_close())
+        del clients[removed]
+
+        if removed in log_files:
+            log_files[removed].close()
+            del log_files[removed]
+
+    for added in TEST_ROOM_IDS - clients.keys():
+        client = None
+        try:
+            client = asyncio.run(connect_room(added))
+        except Exception as e:
+            if client is not None:
+                asyncio.run(client.stop_and_close())
 
 
 async def main():
+    global TEST_ROOM_IDS
+
     flush_files()
 
+    redis.on_room_changed = on_room_changed
+    redis.init(conf.get("redis", None))
+    TEST_ROOM_IDS = list(set(TEST_ROOM_IDS + redis.get_room_list()))
+
     init_session()
     try:
         for room_id in TEST_ROOM_IDS:
@@ -82,9 +81,6 @@ async def main():
 
                 raise e
 
-        # run subscribe_redis on new thread
-        threading.Thread(target=subscribe_redis, daemon=True).start()
-
         while True:
             await asyncio.gather(*(
                 clients[room_id].join() for room_id in clients
@@ -135,35 +131,6 @@ def flush_files():
     timer.start()
 
 
-def subscribe_redis():
-    p = r.pubsub()
-    p.psubscribe(**{f"__keyspace@0__:{room_list_key}": room_changed_handler})
-    for message in p.listen():
-        pass
-
-
-# noinspection PyUnusedLocal
-def room_changed_handler(message):
-    global TEST_ROOM_IDS
-    TEST_ROOM_IDS = set(int(room_id) for room_id in set(r.lrange(room_list_key, 0, -1)))
-
-    for removed in clients.keys() - TEST_ROOM_IDS:
-        asyncio.run(clients[removed].stop_and_close())
-        del clients[removed]
-
-        if removed in log_files:
-            log_files[removed].close()
-            del log_files[removed]
-
-    for added in TEST_ROOM_IDS - clients.keys():
-        client = None
-        try:
-            client = asyncio.run(connect_room(added))
-        except Exception as e:
-            if client is not None:
-                asyncio.run(client.stop_and_close())
-
-
 def init_session():
     cookies = http.cookies.SimpleCookie()
     cookies["SESSDATA"] = SESSDATA
@@ -185,6 +152,33 @@ class JsonlLogger(blivedm.BaseHandler):
         log_files[client.room_id].write(f"{json.dumps(message, separators=(',', ':'), ensure_ascii=False)}\n")
 
 
+# set log level to INFO
+logging.basicConfig(level=logging.INFO)
+
+try:
+    with open("config.json") as conf_f:
+        conf = json.load(conf_f)
+except FileNotFoundError:
+    conf = {}
+redis_conf: dict | None = conf.get("redis", None)
+
+os.makedirs("logs", exist_ok=True)
+room_status_log = open("logs/room_status.jsonl", "a", encoding="utf-8")
+
+logger = logging.getLogger()
+
+clients = {}
+log_files = {}
+
+with open("cookie.txt", "r", encoding="utf-8") as f:
+    SESSDATA = f.read().strip()
+
+session: Optional[aiohttp.ClientSession] = None
+
+TEST_ROOM_IDS: list[int] = room_from_config()
+
+signal.signal(signal.SIGHUP, handle_sighup)
+
 jsonl_logger = JsonlLogger()
 
 if __name__ == "__main__":

redis_util.py

@@ -0,0 +1,43 @@
+import threading
+from typing import Callable
+
+import redis
+
+room_list_key = "bilibili.live.danmu.room_list"
+r: redis.Redis | None = None
+on_room_changed: Callable[[list[int]], None] = lambda ignore: None
+
+
+def init(redis_conf: dict | None):
+    global r
+
+    if redis_conf is None:
+        return
+
+    r = redis.Redis(
+        host=redis_conf.get("host", "127.0.0.1"),
+        port=redis_conf.get("port", 6379),
+        db=redis_conf.get("port", 0),
+    )
+
+    # run subscribe_redis on new thread
+    threading.Thread(target=subscribe_redis, daemon=True).start()
+
+
+def get_room_list() -> list[int]:
+    if r is None:
+        return []
+
+    return [int(room_id) for room_id in set(r.lrange(room_list_key, 0, -1))]
+
+
+def subscribe_redis():
+    p = r.pubsub()
+    p.psubscribe(**{f"__keyspace@0__:{room_list_key}": room_changed_handler})
+    for message in p.listen():
+        pass
+
+
+# noinspection PyUnusedLocal
+def room_changed_handler(ignore):
+    on_room_changed(get_room_list())