Dedupe get_attribute_converter

beetsplug/aura.py

@@ -20,7 +20,7 @@
 from dataclasses import dataclass
 from mimetypes import guess_type
 from os.path import getsize, isfile
-from typing import Mapping
+from typing import ClassVar, Mapping, Type
 
 from flask import (
     Blueprint,
@@ -41,8 +41,9 @@
     NotQuery,
     RegexpQuery,
     SlowFieldSort,
+    SQLiteType,
 )
-from beets.library import Album, Item, Library
+from beets.library import Album, Item, LibModel, Library
 from beets.plugins import BeetsPlugin
 from beets.ui import Subcommand, _open_library
 from beets.util import py3_path
@@ -124,6 +125,8 @@
 class AURADocument:
     """Base class for building AURA documents."""
 
+    model_cls: ClassVar[Type[LibModel]]
+
     lib: Library
     args: Mapping[str, str]
 
@@ -147,6 +150,22 @@
         }
         return make_response(document, status)
 
+    @classmethod
+    def get_attribute_converter(cls, beets_attr: str) -> Type[SQLiteType]:
+        """Work out what data type an attribute should be for beets.
+
+        Args:
+            beets_attr: The name of the beets attribute, e.g. "title".
+        """
+        try:
+            # Look for field in list of Album fields
+            # and get python type of database type.
+            # See beets.library.Album and beets.dbcore.types
+            return cls.model_cls._fields[beets_attr].model_type
+        except KeyError:
+            # Fall back to string (NOTE: probably not good)
+            return str
+
     def translate_filters(self):
         """Translate filters from request arguments to a beets Query."""
         # The format of each filter key in the request parameter is:
@@ -339,6 +358,8 @@
 class TrackDocument(AURADocument):
     """Class for building documents for /tracks endpoints."""
 
+    model_cls = Item
+
     attribute_map = TRACK_ATTR_MAP
 
     def get_collection(self, query=None, sort=None):
@@ -350,29 +371,22 @@
         """
         return self.lib.items(query, sort)
 
-    def get_attribute_converter(self, beets_attr):
+    @classmethod
+    def get_attribute_converter(cls, beets_attr: str) -> Type[SQLiteType]:
         """Work out what data type an attribute should be for beets.
 
         Args:
             beets_attr: The name of the beets attribute, e.g. "title".
         """
         # filesize is a special field (read from disk not db?)
         if beets_attr == "filesize":
-            converter = int
-        else:
-            try:
-                # Look for field in list of Item fields
-                # and get python type of database type.
-                # See beets.library.Item and beets.dbcore.types
-                converter = Item._fields[beets_attr].model_type
-            except KeyError:
-                # Fall back to string (NOTE: probably not good)
-                converter = str
-        return converter
+            return int
+
+        return super().get_attribute_converter(beets_attr)
 
     @staticmethod
     def get_resource_object(lib: Library, track):
         """Construct a JSON:API resource object from a beets Item.
 
         Args:
             track: A beets Item object.
@@ -426,6 +440,8 @@
 class AlbumDocument(AURADocument):
     """Class for building documents for /albums endpoints."""
 
+    model_cls = Album
+
     attribute_map = ALBUM_ATTR_MAP
 
     def get_collection(self, query=None, sort=None):
@@ -437,25 +453,9 @@
         """
         return self.lib.albums(query, sort)
 
-    def get_attribute_converter(self, beets_attr):
-        """Work out what data type an attribute should be for beets.
-
-        Args:
-            beets_attr: The name of the beets attribute, e.g. "title".
-        """
-        try:
-            # Look for field in list of Album fields
-            # and get python type of database type.
-            # See beets.library.Album and beets.dbcore.types
-            converter = Album._fields[beets_attr].model_type
-        except KeyError:
-            # Fall back to string (NOTE: probably not good)
-            converter = str
-        return converter
-
     @staticmethod
     def get_resource_object(lib: Library, album):
         """Construct a JSON:API resource object from a beets Album.
 
         Args:
             album: A beets Album object.
@@ -526,6 +526,8 @@
 class ArtistDocument(AURADocument):
     """Class for building documents for /artists endpoints."""
 
+    model_cls = Item
+
     attribute_map = ARTIST_ATTR_MAP
 
     def get_collection(self, query=None, sort=None):
@@ -544,25 +546,9 @@
                 collection.append(track.artist)
         return collection
 
-    def get_attribute_converter(self, beets_attr):
-        """Work out what data type an attribute should be for beets.
-
-        Args:
-            beets_attr: The name of the beets attribute, e.g. "artist".
-        """
-        try:
-            # Look for field in list of Item fields
-            # and get python type of database type.
-            # See beets.library.Item and beets.dbcore.types
-            converter = Item._fields[beets_attr].model_type
-        except KeyError:
-            # Fall back to string (NOTE: probably not good)
-            converter = str
-        return converter
-
     @staticmethod
     def get_resource_object(lib: Library, artist_id):
         """Construct a JSON:API resource object for the given artist.
 
         Args:
             artist_id: A string which is the artist's name.
@@ -643,9 +629,11 @@
 class ImageDocument(AURADocument):
     """Class for building documents for /images/(id) endpoints."""
 
+    model_cls = Album
+
     @staticmethod
     def get_image_path(lib: Library, image_id):
         """Works out the full path to the image with the given id.
 
         Returns None if there is no such image.
 
@@ -687,7 +675,7 @@
 
     @staticmethod
     def get_resource_object(lib: Library, image_id):
         """Construct a JSON:API resource object for the given image.
 
         Args:
             image_id: A string in the form