diff --git a/blivedm/clients/web.py b/blivedm/clients/web.py index e7a6d78..d9de9d5 100644 --- a/blivedm/clients/web.py +++ b/blivedm/clients/web.py @@ -35,12 +35,12 @@ class BLiveClient(ws_base.WebSocketClientBase): """ def __init__( - self, - room_id: int, - *, - uid: Optional[int] = None, - session: Optional[aiohttp.ClientSession] = None, - heartbeat_interval=30, + self, + room_id: int, + *, + uid: Optional[int] = None, + session: Optional[aiohttp.ClientSession] = None, + heartbeat_interval=30, ): super().__init__(session, heartbeat_interval) @@ -50,6 +50,8 @@ class BLiveClient(ws_base.WebSocketClientBase): # 在调用init_room后初始化的字段 self._room_owner_uid: Optional[int] = None + self.live_status: Optional[int] = None + self.live_start_time: Optional[int] = None """主播用户ID""" self._host_server_list: Optional[List[dict]] = None """ @@ -120,8 +122,8 @@ class BLiveClient(ws_base.WebSocketClientBase): try: async with self._session.get( - UID_INIT_URL, - headers={'User-Agent': utils.USER_AGENT}, + UID_INIT_URL, + headers={'User-Agent': utils.USER_AGENT}, ) as res: if res.status != 200: logger.warning('room=%d _init_uid() failed, status=%d, reason=%s', self._tmp_room_id, @@ -158,8 +160,8 @@ class BLiveClient(ws_base.WebSocketClientBase): async def _init_buvid(self): try: async with self._session.get( - BUVID_INIT_URL, - headers={'User-Agent': utils.USER_AGENT}, + BUVID_INIT_URL, + headers={'User-Agent': utils.USER_AGENT}, ) as res: if res.status != 200: logger.warning('room=%d _init_buvid() status error, status=%d, reason=%s', @@ -171,11 +173,11 @@ class BLiveClient(ws_base.WebSocketClientBase): async def _init_room_id_and_owner(self): try: async with self._session.get( - ROOM_INIT_URL, - headers={'User-Agent': utils.USER_AGENT}, - params={ - 'room_id': self._tmp_room_id - }, + ROOM_INIT_URL, + headers={'User-Agent': utils.USER_AGENT}, + params={ + 'room_id': self._tmp_room_id + }, ) as res: if res.status != 200: logger.warning('room=%d _init_room_id_and_owner() failed, status=%d, reason=%s', self._tmp_room_id, @@ -197,17 +199,19 @@ class BLiveClient(ws_base.WebSocketClientBase): room_info = data['room_info'] self._room_id = room_info['room_id'] self._room_owner_uid = room_info['uid'] + self.live_status = room_info['live_status'] + self.live_start_time = room_info['live_start_time'] return True async def _init_host_server(self): try: async with self._session.get( - DANMAKU_SERVER_CONF_URL, - headers={'User-Agent': utils.USER_AGENT}, - params={ - 'id': self._room_id, - 'type': 0 - }, + DANMAKU_SERVER_CONF_URL, + headers={'User-Agent': utils.USER_AGENT}, + params={ + 'id': self._room_id, + 'type': 0 + }, ) as res: if res.status != 200: logger.warning('room=%d _init_host_server() failed, status=%d, reason=%s', self._room_id, diff --git a/sample.py b/sample.py index cea9bbd..6661e01 100644 --- a/sample.py +++ b/sample.py @@ -13,13 +13,16 @@ import redis import blivedm from blivedm.clients import ws_base +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', port=6379, db=0) +r = redis.Redis(host="redis", port=6379, db=0) TEST_ROOM_IDS = [int(room_id) for room_id in set(r.lrange(room_list_key, 0, -1))] clients = {} @@ -39,12 +42,7 @@ async def main(): for room_id in TEST_ROOM_IDS: client = None try: - logger.info("connect to room %s", room_id) - client = blivedm.BLiveClient(room_id, session=session) - client.set_handler(jsonl_logger) - client.start() - clients[room_id] = client - logger.info("connect to room %s success", room_id) + client = await connect_room(room_id) except Exception as e: if client is not None: await client.stop_and_close() @@ -65,7 +63,32 @@ async def main(): await session.close() +async def connect_room(room_id): + logger.info("connect to room %s", room_id) + client = blivedm.BLiveClient(room_id, session=session) + client.set_handler(jsonl_logger) + await client.init_room() + client._need_init_room = False + client.start() + clients[room_id] = client + logger.info("connect to room %s success", room_id) + + room_status_log.write(json.dumps( + { + "room_id": room_id, + "live_status": client.live_status, + "live_start_time": client.live_start_time, + }, + separators=(",", ":"), + ensure_ascii=False, + )) + room_status_log.write("\n") + return client + + def flush_files(): + room_status_log.flush() + for room_id in log_files: log_files[room_id].flush() @@ -97,12 +120,7 @@ def room_changed_handler(message): for added in TEST_ROOM_IDS - clients.keys(): client = None try: - logger.info("connect to room %s", added) - client = blivedm.BLiveClient(added, session=session) - client.set_handler(jsonl_logger) - client.start() - clients[added] = client - logger.info("connect to room %s success", added) + client = asyncio.run(connect_room(added)) except Exception as e: if client is not None: asyncio.run(client.stop_and_close()) @@ -112,8 +130,8 @@ def room_changed_handler(message): def init_session(): cookies = http.cookies.SimpleCookie() - cookies['SESSDATA'] = SESSDATA - cookies['SESSDATA']['domain'] = 'bilibili.com' + cookies["SESSDATA"] = SESSDATA + cookies["SESSDATA"]["domain"] = "bilibili.com" global session session = aiohttp.ClientSession() @@ -123,7 +141,6 @@ def init_session(): class JsonlLogger(blivedm.BaseHandler): def _global_callback(self, client: ws_base.WebSocketClientBase, message: dict): if client.room_id not in log_files: - os.makedirs("logs", exist_ok=True) log_files[client.room_id] = open(f"logs/{client.room_id}.jsonl", "a", encoding="utf-8") log_files[client.room_id].write(f"{json.dumps(message, separators=(',', ':'), ensure_ascii=False)}\n") @@ -131,5 +148,5 @@ class JsonlLogger(blivedm.BaseHandler): jsonl_logger = JsonlLogger() -if __name__ == '__main__': +if __name__ == "__main__": asyncio.run(main())