Dedupe get_attribute_converter

Author: snejus

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,
@@ -42,7 +42,7 @@
     RegexpQuery,
     SlowFieldSort,
 )
-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 +124,8 @@
 class AURADocument:
     """Base class for building AURA documents."""
 
+    model_cls: ClassVar[Type[LibModel]]
+
     lib: Library
     args: Mapping[str, str]
 
@@ -147,6 +149,22 @@ def error(status, title, detail):
         }
         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 +357,8 @@ def single_resource_document(self, resource_object):
 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,25 +370,18 @@ def get_collection(self, query=None, sort=None):
         """
         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):
@@ -426,6 +439,8 @@ def single_resource(self, track_id):
 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,22 +452,6 @@ def get_collection(self, query=None, sort=None):
         """
         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.
@@ -526,6 +525,8 @@ def single_resource(self, album_id):
 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,22 +545,6 @@ def get_collection(self, query=None, sort=None):
                 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.
@@ -643,6 +628,8 @@ def safe_filename(fn):
 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.