Skip to content

Commit d02dafd

Browse files
committed
Dedupe get_attribute_converter
1 parent 78d381c commit d02dafd

File tree

1 file changed

+34
-46
lines changed

1 file changed

+34
-46
lines changed

beetsplug/aura.py

+34-46
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@
2020
from dataclasses import dataclass
2121
from mimetypes import guess_type
2222
from os.path import getsize, isfile
23-
from typing import Mapping
23+
from typing import ClassVar, Mapping, Type
2424

2525
from flask import (
2626
Blueprint,
@@ -41,8 +41,9 @@
4141
NotQuery,
4242
RegexpQuery,
4343
SlowFieldSort,
44+
SQLiteType,
4445
)
45-
from beets.library import Album, Item, Library
46+
from beets.library import Album, Item, LibModel, Library
4647
from beets.plugins import BeetsPlugin
4748
from beets.ui import Subcommand, _open_library
4849
from beets.util import py3_path
@@ -124,6 +125,8 @@
124125
class AURADocument:
125126
"""Base class for building AURA documents."""
126127

128+
model_cls: ClassVar[Type[LibModel]]
129+
127130
lib: Library
128131
args: Mapping[str, str]
129132

@@ -147,6 +150,22 @@ def error(status, title, detail):
147150
}
148151
return make_response(document, status)
149152

153+
@classmethod
154+
def get_attribute_converter(cls, beets_attr: str) -> Type[SQLiteType]:
155+
"""Work out what data type an attribute should be for beets.
156+
157+
Args:
158+
beets_attr: The name of the beets attribute, e.g. "title".
159+
"""
160+
try:
161+
# Look for field in list of Album fields
162+
# and get python type of database type.
163+
# See beets.library.Album and beets.dbcore.types
164+
return cls.model_cls._fields[beets_attr].model_type
165+
except KeyError:
166+
# Fall back to string (NOTE: probably not good)
167+
return str
168+
150169
def translate_filters(self):
151170
"""Translate filters from request arguments to a beets Query."""
152171
# The format of each filter key in the request parameter is:
@@ -339,6 +358,8 @@ def single_resource_document(self, resource_object):
339358
class TrackDocument(AURADocument):
340359
"""Class for building documents for /tracks endpoints."""
341360

361+
model_cls = Item
362+
342363
attribute_map = TRACK_ATTR_MAP
343364

344365
def get_collection(self, query=None, sort=None):
@@ -350,25 +371,18 @@ def get_collection(self, query=None, sort=None):
350371
"""
351372
return self.lib.items(query, sort)
352373

353-
def get_attribute_converter(self, beets_attr):
374+
@classmethod
375+
def get_attribute_converter(cls, beets_attr: str) -> Type[SQLiteType]:
354376
"""Work out what data type an attribute should be for beets.
355377
356378
Args:
357379
beets_attr: The name of the beets attribute, e.g. "title".
358380
"""
359381
# filesize is a special field (read from disk not db?)
360382
if beets_attr == "filesize":
361-
converter = int
362-
else:
363-
try:
364-
# Look for field in list of Item fields
365-
# and get python type of database type.
366-
# See beets.library.Item and beets.dbcore.types
367-
converter = Item._fields[beets_attr].model_type
368-
except KeyError:
369-
# Fall back to string (NOTE: probably not good)
370-
converter = str
371-
return converter
383+
return int
384+
385+
return super().get_attribute_converter(beets_attr)
372386

373387
@staticmethod
374388
def get_resource_object(lib: Library, track):
@@ -426,6 +440,8 @@ def single_resource(self, track_id):
426440
class AlbumDocument(AURADocument):
427441
"""Class for building documents for /albums endpoints."""
428442

443+
model_cls = Album
444+
429445
attribute_map = ALBUM_ATTR_MAP
430446

431447
def get_collection(self, query=None, sort=None):
@@ -437,22 +453,6 @@ def get_collection(self, query=None, sort=None):
437453
"""
438454
return self.lib.albums(query, sort)
439455

440-
def get_attribute_converter(self, beets_attr):
441-
"""Work out what data type an attribute should be for beets.
442-
443-
Args:
444-
beets_attr: The name of the beets attribute, e.g. "title".
445-
"""
446-
try:
447-
# Look for field in list of Album fields
448-
# and get python type of database type.
449-
# See beets.library.Album and beets.dbcore.types
450-
converter = Album._fields[beets_attr].model_type
451-
except KeyError:
452-
# Fall back to string (NOTE: probably not good)
453-
converter = str
454-
return converter
455-
456456
@staticmethod
457457
def get_resource_object(lib: Library, album):
458458
"""Construct a JSON:API resource object from a beets Album.
@@ -526,6 +526,8 @@ def single_resource(self, album_id):
526526
class ArtistDocument(AURADocument):
527527
"""Class for building documents for /artists endpoints."""
528528

529+
model_cls = Item
530+
529531
attribute_map = ARTIST_ATTR_MAP
530532

531533
def get_collection(self, query=None, sort=None):
@@ -544,22 +546,6 @@ def get_collection(self, query=None, sort=None):
544546
collection.append(track.artist)
545547
return collection
546548

547-
def get_attribute_converter(self, beets_attr):
548-
"""Work out what data type an attribute should be for beets.
549-
550-
Args:
551-
beets_attr: The name of the beets attribute, e.g. "artist".
552-
"""
553-
try:
554-
# Look for field in list of Item fields
555-
# and get python type of database type.
556-
# See beets.library.Item and beets.dbcore.types
557-
converter = Item._fields[beets_attr].model_type
558-
except KeyError:
559-
# Fall back to string (NOTE: probably not good)
560-
converter = str
561-
return converter
562-
563549
@staticmethod
564550
def get_resource_object(lib: Library, artist_id):
565551
"""Construct a JSON:API resource object for the given artist.
@@ -643,6 +629,8 @@ def safe_filename(fn):
643629
class ImageDocument(AURADocument):
644630
"""Class for building documents for /images/(id) endpoints."""
645631

632+
model_cls = Album
633+
646634
@staticmethod
647635
def get_image_path(lib: Library, image_id):
648636
"""Works out the full path to the image with the given id.

0 commit comments

Comments
 (0)