Move MusicBrainzAPI to a shared util

This commit is contained in:
Šarūnas Nejus 2025-12-22 04:04:02 +00:00
parent ea2e7bf997
commit 523fa6ceaf
No known key found for this signature in database
6 changed files with 229 additions and 204 deletions

View file

@ -0,0 +1,122 @@
from __future__ import annotations
import operator
from dataclasses import dataclass, field
from functools import cached_property, singledispatchmethod
from itertools import groupby
from typing import TYPE_CHECKING, Any
from requests_ratelimiter import LimiterMixin
from beets import config
from .requests import RequestHandler, TimeoutAndRetrySession
if TYPE_CHECKING:
from .._typing import JSONDict
class LimiterTimeoutSession(LimiterMixin, TimeoutAndRetrySession):
pass
@dataclass
class MusicBrainzAPI(RequestHandler):
api_host: str = field(init=False)
rate_limit: float = field(init=False)
def __post_init__(self) -> None:
mb_config = config["musicbrainz"]
mb_config.add(
{
"host": "musicbrainz.org",
"https": False,
"ratelimit": 1,
"ratelimit_interval": 1,
}
)
hostname = mb_config["host"].as_str()
if hostname == "musicbrainz.org":
self.api_host, self.rate_limit = "https://musicbrainz.org", 1.0
else:
https = mb_config["https"].get(bool)
self.api_host = f"http{'s' if https else ''}://{hostname}"
self.rate_limit = (
mb_config["ratelimit"].get(int)
/ mb_config["ratelimit_interval"].as_number()
)
def create_session(self) -> LimiterTimeoutSession:
return LimiterTimeoutSession(per_second=self.rate_limit)
def get_entity(
self, entity: str, includes: list[str] | None = None, **kwargs
) -> JSONDict:
if includes:
kwargs["inc"] = "+".join(includes)
return self._group_relations(
self.get_json(
f"{self.api_host}/ws/2/{entity}",
params={**kwargs, "fmt": "json"},
)
)
def get_release(self, id_: str, **kwargs) -> JSONDict:
return self.get_entity(f"release/{id_}", **kwargs)
def get_recording(self, id_: str, **kwargs) -> JSONDict:
return self.get_entity(f"recording/{id_}", **kwargs)
def browse_recordings(self, **kwargs) -> list[JSONDict]:
return self.get_entity("recording", **kwargs)["recordings"]
@singledispatchmethod
@classmethod
def _group_relations(cls, data: Any) -> Any:
"""Normalize MusicBrainz 'relations' into type-keyed fields recursively.
This helper rewrites payloads that use a generic 'relations' list into
a structure that is easier to consume downstream. When a mapping
contains 'relations', those entries are regrouped by their 'target-type'
and stored under keys like '<target-type>-relations'. The original
'relations' key is removed to avoid ambiguous access patterns.
The transformation is applied recursively so that nested objects and
sequences are normalized consistently, while non-container values are
left unchanged.
"""
return data
@_group_relations.register(list)
@classmethod
def _(cls, data: list[Any]) -> list[Any]:
return [cls._group_relations(i) for i in data]
@_group_relations.register(dict)
@classmethod
def _(cls, data: JSONDict) -> JSONDict:
for k, v in list(data.items()):
if k == "relations":
get_target_type = operator.methodcaller("get", "target-type")
for target_type, group in groupby(
sorted(v, key=get_target_type), get_target_type
):
relations = [
{k: v for k, v in item.items() if k != "target-type"}
for item in group
]
data[f"{target_type}-relations"] = cls._group_relations(
relations
)
data.pop("relations")
else:
data[k] = cls._group_relations(v)
return data
class MusicBrainzAPIMixin:
@cached_property
def mb_api(self) -> MusicBrainzAPI:
return MusicBrainzAPI()

View file

@ -141,7 +141,7 @@ class MusicBrainzPseudoReleasePlugin(MusicBrainzPlugin):
if (ids := self._intercept_mb_release(release)) and (
album_id := self._extract_id(ids[0])
):
raw_pseudo_release = self.api.get_release(album_id)
raw_pseudo_release = self.mb_api.get_release(album_id)
pseudo_release = super().album_info(raw_pseudo_release)
if self.config["custom_tags_only"].get(bool):

View file

@ -16,17 +16,14 @@
from __future__ import annotations
import operator
from collections import Counter
from contextlib import suppress
from dataclasses import dataclass
from functools import cached_property, singledispatchmethod
from itertools import groupby, product
from functools import cached_property
from itertools import product
from typing import TYPE_CHECKING, Any
from urllib.parse import urljoin
from confuse.exceptions import NotFoundError
from requests_ratelimiter import LimiterMixin
import beets
import beets.autotag.hooks
@ -35,11 +32,8 @@ from beets.metadata_plugins import MetadataSourcePlugin
from beets.util.deprecation import deprecate_for_user
from beets.util.id_extractors import extract_release_id
from ._utils.requests import (
HTTPNotFoundError,
RequestHandler,
TimeoutAndRetrySession,
)
from ._utils.musicbrainz import MusicBrainzAPIMixin
from ._utils.requests import HTTPNotFoundError
if TYPE_CHECKING:
from collections.abc import Iterable, Sequence
@ -103,86 +97,6 @@ BROWSE_CHUNKSIZE = 100
BROWSE_MAXTRACKS = 500
class LimiterTimeoutSession(LimiterMixin, TimeoutAndRetrySession):
pass
@dataclass
class MusicBrainzAPI(RequestHandler):
api_host: str
rate_limit: float
def create_session(self) -> LimiterTimeoutSession:
return LimiterTimeoutSession(per_second=self.rate_limit)
def get_entity(
self, entity: str, inc_list: list[str] | None = None, **kwargs
) -> JSONDict:
if inc_list:
kwargs["inc"] = "+".join(inc_list)
return self._group_relations(
self.get_json(
f"{self.api_host}/ws/2/{entity}",
params={**kwargs, "fmt": "json"},
)
)
def get_release(self, id_: str) -> JSONDict:
return self.get_entity(f"release/{id_}", inc_list=RELEASE_INCLUDES)
def get_recording(self, id_: str) -> JSONDict:
return self.get_entity(f"recording/{id_}", inc_list=TRACK_INCLUDES)
def browse_recordings(self, **kwargs) -> list[JSONDict]:
kwargs.setdefault("limit", BROWSE_CHUNKSIZE)
kwargs.setdefault("inc_list", BROWSE_INCLUDES)
return self.get_entity("recording", **kwargs)["recordings"]
@singledispatchmethod
@classmethod
def _group_relations(cls, data: Any) -> Any:
"""Normalize MusicBrainz 'relations' into type-keyed fields recursively.
This helper rewrites payloads that use a generic 'relations' list into
a structure that is easier to consume downstream. When a mapping
contains 'relations', those entries are regrouped by their 'target-type'
and stored under keys like '<target-type>-relations'. The original
'relations' key is removed to avoid ambiguous access patterns.
The transformation is applied recursively so that nested objects and
sequences are normalized consistently, while non-container values are
left unchanged.
"""
return data
@_group_relations.register(list)
@classmethod
def _(cls, data: list[Any]) -> list[Any]:
return [cls._group_relations(i) for i in data]
@_group_relations.register(dict)
@classmethod
def _(cls, data: JSONDict) -> JSONDict:
for k, v in list(data.items()):
if k == "relations":
get_target_type = operator.methodcaller("get", "target-type")
for target_type, group in groupby(
sorted(v, key=get_target_type), get_target_type
):
relations = [
{k: v for k, v in item.items() if k != "target-type"}
for item in group
]
data[f"{target_type}-relations"] = cls._group_relations(
relations
)
data.pop("relations")
else:
data[k] = cls._group_relations(v)
return data
def _preferred_alias(
aliases: list[JSONDict], languages: list[str] | None = None
) -> JSONDict | None:
@ -405,25 +319,11 @@ def _merge_pseudo_and_actual_album(
return merged
class MusicBrainzPlugin(MetadataSourcePlugin):
class MusicBrainzPlugin(MusicBrainzAPIMixin, MetadataSourcePlugin):
@cached_property
def genres_field(self) -> str:
return f"{self.config['genres_tag'].as_choice(['genre', 'tag'])}s"
@cached_property
def api(self) -> MusicBrainzAPI:
hostname = self.config["host"].as_str()
if hostname == "musicbrainz.org":
hostname, rate_limit = "https://musicbrainz.org", 1.0
else:
https = self.config["https"].get(bool)
hostname = f"http{'s' if https else ''}://{hostname}"
rate_limit = (
self.config["ratelimit"].get(int)
/ self.config["ratelimit_interval"].as_number()
)
return MusicBrainzAPI(hostname, rate_limit)
def __init__(self):
"""Set up the python-musicbrainz-ngs module according to settings
from the beets configuration. This should be called at startup.
@ -431,10 +331,6 @@ class MusicBrainzPlugin(MetadataSourcePlugin):
super().__init__()
self.config.add(
{
"host": "musicbrainz.org",
"https": False,
"ratelimit": 1,
"ratelimit_interval": 1,
"genres": False,
"genres_tag": "genre",
"external_ids": {
@ -589,7 +485,9 @@ class MusicBrainzPlugin(MetadataSourcePlugin):
for i in range(0, ntracks, BROWSE_CHUNKSIZE):
self._log.debug("Retrieving tracks starting at {}", i)
recording_list.extend(
self.api.browse_recordings(release=release["id"], offset=i)
self.mb_api.browse_recordings(
release=release["id"], offset=i
)
)
track_map = {r["id"]: r for r in recording_list}
for medium in release["media"]:
@ -861,7 +759,7 @@ class MusicBrainzPlugin(MetadataSourcePlugin):
self._log.debug(
"Searching for MusicBrainz {}s with: {!r}", query_type, query
)
return self.api.get_entity(
return self.mb_api.get_entity(
query_type, query=query, limit=self.config["search_limit"].get()
)[f"{query_type}s"]
@ -901,7 +799,7 @@ class MusicBrainzPlugin(MetadataSourcePlugin):
self._log.debug("Invalid MBID ({}).", album_id)
return None
res = self.api.get_release(albumid)
res = self.mb_api.get_release(albumid, includes=RELEASE_INCLUDES)
# resolve linked release relations
actual_res = None
@ -914,7 +812,9 @@ class MusicBrainzPlugin(MetadataSourcePlugin):
rel["type"] == "transl-tracklisting"
and rel["direction"] == "backward"
):
actual_res = self.api.get_release(rel["release"]["id"])
actual_res = self.mb_api.get_release(
rel["release"]["id"], includes=RELEASE_INCLUDES
)
# release is potentially a pseudo release
release = self.album_info(res)
@ -937,6 +837,8 @@ class MusicBrainzPlugin(MetadataSourcePlugin):
return None
with suppress(HTTPNotFoundError):
return self.track_info(self.api.get_recording(trackid))
return self.track_info(
self.mb_api.get_recording(trackid, includes=TRACK_INCLUDES)
)
return None

View file

@ -94,7 +94,7 @@ class TestMBPseudoMixin(PluginMixin):
@pytest.fixture(autouse=True)
def patch_get_release(self, monkeypatch, pseudo_release: JSONDict):
monkeypatch.setattr(
"beetsplug.musicbrainz.MusicBrainzAPI.get_release",
"beetsplug._utils.musicbrainz.MusicBrainzAPI.get_release",
lambda _, album_id: deepcopy(
{pseudo_release["id"]: pseudo_release}[album_id]
),

View file

@ -863,7 +863,7 @@ class MBLibraryTest(MusicBrainzTestCase):
]
with mock.patch(
"beetsplug.musicbrainz.MusicBrainzAPI.get_release"
"beetsplug._utils.musicbrainz.MusicBrainzAPI.get_release"
) as gp:
gp.side_effect = side_effect
album = self.mb.album_for_id("d2a6f856-b553-40a0-ac54-a321e8e2da02")
@ -907,7 +907,7 @@ class MBLibraryTest(MusicBrainzTestCase):
]
with mock.patch(
"beetsplug.musicbrainz.MusicBrainzAPI.get_release"
"beetsplug._utils.musicbrainz.MusicBrainzAPI.get_release"
) as gp:
gp.side_effect = side_effect
album = self.mb.album_for_id("d2a6f856-b553-40a0-ac54-a321e8e2da02")
@ -951,7 +951,7 @@ class MBLibraryTest(MusicBrainzTestCase):
]
with mock.patch(
"beetsplug.musicbrainz.MusicBrainzAPI.get_release"
"beetsplug._utils.musicbrainz.MusicBrainzAPI.get_release"
) as gp:
gp.side_effect = side_effect
album = self.mb.album_for_id("d2a6f856-b553-40a0-ac54-a321e8e2da02")
@ -1004,7 +1004,7 @@ class MBLibraryTest(MusicBrainzTestCase):
]
with mock.patch(
"beetsplug.musicbrainz.MusicBrainzAPI.get_release"
"beetsplug._utils.musicbrainz.MusicBrainzAPI.get_release"
) as gp:
gp.side_effect = side_effect
album = self.mb.album_for_id("d2a6f856-b553-40a0-ac54-a321e8e2da02")
@ -1055,7 +1055,7 @@ class TestMusicBrainzPlugin(PluginMixin):
def test_item_candidates(self, monkeypatch, mb):
monkeypatch.setattr(
"beetsplug.musicbrainz.MusicBrainzAPI.get_json",
"beetsplug._utils.musicbrainz.MusicBrainzAPI.get_json",
lambda *_, **__: {"recordings": [self.RECORDING]},
)
@ -1066,11 +1066,11 @@ class TestMusicBrainzPlugin(PluginMixin):
def test_candidates(self, monkeypatch, mb):
monkeypatch.setattr(
"beetsplug.musicbrainz.MusicBrainzAPI.get_json",
"beetsplug._utils.musicbrainz.MusicBrainzAPI.get_json",
lambda *_, **__: {"releases": [{"id": self.mbid}]},
)
monkeypatch.setattr(
"beetsplug.musicbrainz.MusicBrainzAPI.get_release",
"beetsplug._utils.musicbrainz.MusicBrainzAPI.get_release",
lambda *_, **__: {
"title": "hi",
"id": self.mbid,
@ -1099,84 +1099,3 @@ class TestMusicBrainzPlugin(PluginMixin):
assert len(candidates) == 1
assert candidates[0].tracks[0].track_id == self.RECORDING["id"]
assert candidates[0].album == "hi"
def test_group_relations():
raw_release = {
"id": "r1",
"relations": [
{"target-type": "artist", "type": "vocal", "name": "A"},
{"target-type": "url", "type": "streaming", "url": "http://s"},
{"target-type": "url", "type": "purchase", "url": "http://p"},
{
"target-type": "work",
"type": "performance",
"work": {
"relations": [
{
"artist": {"name": "幾田りら"},
"target-type": "artist",
"type": "composer",
},
{
"target-type": "url",
"type": "lyrics",
"url": {
"resource": "https://utaten.com/lyric/tt24121002/"
},
},
{
"artist": {"name": "幾田りら"},
"target-type": "artist",
"type": "lyricist",
},
{
"target-type": "url",
"type": "lyrics",
"url": {
"resource": "https://www.uta-net.com/song/366579/"
},
},
],
"title": "百花繚乱",
"type": "Song",
},
},
],
}
assert musicbrainz.MusicBrainzAPI._group_relations(raw_release) == {
"id": "r1",
"artist-relations": [{"type": "vocal", "name": "A"}],
"url-relations": [
{"type": "streaming", "url": "http://s"},
{"type": "purchase", "url": "http://p"},
],
"work-relations": [
{
"type": "performance",
"work": {
"artist-relations": [
{"type": "composer", "artist": {"name": "幾田りら"}},
{"type": "lyricist", "artist": {"name": "幾田りら"}},
],
"url-relations": [
{
"type": "lyrics",
"url": {
"resource": "https://utaten.com/lyric/tt24121002/"
},
},
{
"type": "lyrics",
"url": {
"resource": "https://www.uta-net.com/song/366579/"
},
},
],
"title": "百花繚乱",
"type": "Song",
},
},
],
}

View file

@ -0,0 +1,82 @@
from beetsplug._utils.musicbrainz import MusicBrainzAPI
def test_group_relations():
raw_release = {
"id": "r1",
"relations": [
{"target-type": "artist", "type": "vocal", "name": "A"},
{"target-type": "url", "type": "streaming", "url": "http://s"},
{"target-type": "url", "type": "purchase", "url": "http://p"},
{
"target-type": "work",
"type": "performance",
"work": {
"relations": [
{
"artist": {"name": "幾田りら"},
"target-type": "artist",
"type": "composer",
},
{
"target-type": "url",
"type": "lyrics",
"url": {
"resource": "https://utaten.com/lyric/tt24121002/"
},
},
{
"artist": {"name": "幾田りら"},
"target-type": "artist",
"type": "lyricist",
},
{
"target-type": "url",
"type": "lyrics",
"url": {
"resource": "https://www.uta-net.com/song/366579/"
},
},
],
"title": "百花繚乱",
"type": "Song",
},
},
],
}
assert MusicBrainzAPI._group_relations(raw_release) == {
"id": "r1",
"artist-relations": [{"type": "vocal", "name": "A"}],
"url-relations": [
{"type": "streaming", "url": "http://s"},
{"type": "purchase", "url": "http://p"},
],
"work-relations": [
{
"type": "performance",
"work": {
"artist-relations": [
{"type": "composer", "artist": {"name": "幾田りら"}},
{"type": "lyricist", "artist": {"name": "幾田りら"}},
],
"url-relations": [
{
"type": "lyrics",
"url": {
"resource": "https://utaten.com/lyric/tt24121002/"
},
},
{
"type": "lyrics",
"url": {
"resource": "https://www.uta-net.com/song/366579/"
},
},
],
"title": "百花繚乱",
"type": "Song",
},
},
],
}