diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..b52f889 --- /dev/null +++ b/.gitignore @@ -0,0 +1,11 @@ +# Python ignores +__pycache__/ +*.py[cod] +*.so +*.egg-info/ +.eggs/ +.pytest_cache/ +.cache/ + +.openhands/ + diff --git a/net_stream/bilibli_live.py b/net_stream/bilibli_live.py index afdb1b4..9b7c937 100644 --- a/net_stream/bilibli_live.py +++ b/net_stream/bilibli_live.py @@ -1,20 +1,25 @@ import asyncio import requests import numpy as np -from queue import SimpleQueue class BilibiliLive: - def __init__(self, room_id): + def __init__(self, room_id, platform: str = "web", quality: int | None = None, qn: int | None = None, prefer_lowest: bool = True): self.room_id = room_id + # Streaming options + self.platform = platform # 'web' (http-flv) or 'h5' (hls) + self.quality = quality # 2: 流畅, 3: 高清, 4: 原画 (API 'quality' param) + self.qn = qn # 80, 150, 400, 10000, 20000, 30000 (API 'qn' param) + self.prefer_lowest = prefer_lowest + def default_read_audio(): raise RuntimeError("You should call spin_ffmpeg first.") - + self.default_read_audio = default_read_audio self.read_audio = default_read_audio self.process = None self.reader_task = None - + self.ua = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_12_6) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/59.0.3071.115 Safari/537.36" self.stream_play_url = "http://api.live.bilibili.com/room/v1/Room/playUrl" @@ -23,13 +28,55 @@ def get_stream_url(self): "User-Agent": self.ua } - params = { - "cid": self.room_id, + base_params = { + "cid": self.room_id, + "platform": self.platform, } - response = requests.get(self.stream_play_url, params=params, headers=headers) - - return response.json()["data"]["durl"][0]["url"] # TODO: Maybe do durl selection to ensure low bandwidth / low latency + # If explicit quality is provided, request directly. + if self.qn is not None or self.quality is not None: + direct_params = base_params.copy() + if self.qn is not None: + direct_params["qn"] = self.qn + if self.quality is not None: + direct_params["quality"] = self.quality + response = requests.get(self.stream_play_url, params=direct_params, headers=headers) + return response.json()["data"]["durl"][0]["url"] + + # Default behavior: choose the lowest available quality + response = requests.get(self.stream_play_url, params=base_params, headers=headers) + data = response.json().get("data", {}) + + chosen_quality: int | None = None + chosen_qn: int | None = None + accept_quality = data.get("accept_quality") + if isinstance(accept_quality, list) and accept_quality: + try: + codes = [int(x) for x in accept_quality] + m = min(codes) + # Heuristic: small numbers (<=10) correspond to 'quality' codes like 2/3/4, + # otherwise they are 'qn' codes like 80/150/400/10000. + if m <= 10: + chosen_quality = m + else: + chosen_qn = m + except Exception: + chosen_quality = None + chosen_qn = None + + # Fallback if accept_quality is not present or parsing failed + if chosen_quality is None and chosen_qn is None: + # Default to commonly lowest qn is 80 (流畅) + chosen_qn = 80 + + # Request the URL for the chosen lowest quality + direct_params = base_params.copy() + if chosen_quality is not None: + direct_params["quality"] = chosen_quality + else: + direct_params["qn"] = chosen_qn + response2 = requests.get(self.stream_play_url, params=direct_params, headers=headers) + return response2.json()["data"]["durl"][0]["url"] async def spin_ffmpeg(self, ffmpeg_path: str, sampling_rate=16000, samples_per_chunk=512): command = [ @@ -83,6 +130,6 @@ async def stop_ffmpeg(self): if self.reader_task is not None: await self.reader_task self.reader_task = None - + self.read_audio = self.default_read_audio diff --git a/pytest.ini b/pytest.ini new file mode 100644 index 0000000..27eec68 --- /dev/null +++ b/pytest.ini @@ -0,0 +1,3 @@ +[pytest] +testpaths = tests +python_files = test_*.py diff --git a/tests/test_bilibili_live.py b/tests/test_bilibili_live.py new file mode 100644 index 0000000..efe9970 --- /dev/null +++ b/tests/test_bilibili_live.py @@ -0,0 +1,128 @@ +import types +import pytest +import os, sys + +# Ensure project root is on path +ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) +if ROOT not in sys.path: + sys.path.insert(0, ROOT) + +import net_stream.bilibli_live as bl + + +class DummyResponse: + def __init__(self, payload): + self._payload = payload + + def json(self): + return self._payload + + +def test_default_selects_lowest_quality(monkeypatch): + calls = [] + + def fake_get(url, params=None, headers=None): + calls.append({"url": url, "params": dict(params or {}), "headers": dict(headers or {})}) + params = params or {} + # If asking for a concrete quality, return a URL indicating that quality + if "quality" in params: + q = params["quality"] + return DummyResponse({ + "code": 0, + "message": "0", + "ttl": 1, + "data": { + "durl": [{"url": f"http://example.com/stream_q{q}.flv"}] + } + }) + # Initial probe without quality: return accept_quality list (strings) + return DummyResponse({ + "code": 0, + "message": "0", + "ttl": 1, + "data": { + "accept_quality": ["4", "3", "2"], + "durl": [{"url": "http://example.com/base.flv"}], + } + }) + + monkeypatch.setattr(bl.requests, "get", fake_get) + + live = bl.BilibiliLive(room_id=14073662) + url = live.get_stream_url() + + # It should make two requests: probe then fetch with the lowest quality=2 + assert len(calls) == 2 + assert calls[0]["params"].get("cid") == 14073662 + assert calls[0]["params"].get("platform") == "web" + assert "quality" not in calls[0]["params"] and "qn" not in calls[0]["params"] + + assert calls[1]["params"].get("quality") == 2 + assert url.endswith("stream_q2.flv") + + +def test_explicit_quality_is_respected(monkeypatch): + calls = [] + + def fake_get(url, params=None, headers=None): + calls.append({"url": url, "params": dict(params or {}), "headers": dict(headers or {})}) + params = params or {} + if "quality" in params: + q = params["quality"] + return DummyResponse({ + "code": 0, + "message": "0", + "ttl": 1, + "data": { + "durl": [{"url": f"http://example.com/stream_q{q}.flv"}] + } + }) + # Should not reach here in this test + raise AssertionError("Unexpected probe without explicit quality") + + monkeypatch.setattr(bl.requests, "get", fake_get) + + live = bl.BilibiliLive(room_id=1, quality=3) + url = live.get_stream_url() + + assert len(calls) == 1 + assert calls[0]["params"].get("quality") == 3 + assert url.endswith("stream_q3.flv") + + +def test_fallback_uses_qn_80_when_accept_quality_missing(monkeypatch): + calls = [] + + def fake_get(url, params=None, headers=None): + calls.append({"url": url, "params": dict(params or {}), "headers": dict(headers or {})}) + params = params or {} + if "qn" in params: + qn = params["qn"] + return DummyResponse({ + "code": 0, + "message": "0", + "ttl": 1, + "data": { + "durl": [{"url": f"http://example.com/stream_qn{qn}.flv"}] + } + }) + # First probe returns no accept_quality + return DummyResponse({ + "code": 0, + "message": "0", + "ttl": 1, + "data": { + # intentionally missing 'accept_quality' + "durl": [{"url": "http://example.com/base.flv"}], + } + }) + + monkeypatch.setattr(bl.requests, "get", fake_get) + + live = bl.BilibiliLive(room_id=2) + url = live.get_stream_url() + + assert len(calls) == 2 + # Second call should include qn=80 as fallback + assert calls[1]["params"].get("qn") == 80 + assert url.endswith("stream_qn80.flv")