typing: type beets.library.library

beets/library/library.py

@@ -4,7 +4,7 @@ import re
 from contextlib import contextmanager
 from functools import cached_property
 from pathlib import Path
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, TypeVar
 
 import platformdirs
 
@@ -20,10 +20,17 @@ from .models import Album, Item
 from .queries import parse_query_parts, parse_query_string
 
 if TYPE_CHECKING:
-    from beets.dbcore import Results
+    from collections.abc import Iterator, Sequence
+
+    from beets.dbcore.query import Query
+    from beets.dbcore.sort import Sort
     from beets.util import PathLike, Replacements
     from beets.util.pathformats import PathFormat
 
+    from .models import LibModel
+
+    LM = TypeVar("LM", bound=LibModel)
+
 
 class Library(dbcore.Database):
     """A database of music containing songs and albums."""
@@ -65,7 +72,7 @@ class Library(dbcore.Database):
         path: PathLike = Path("library.blb"),
         directory: str | None = None,
         set_music_dir: bool = True,
-    ):
+    ) -> None:
         self.directory = normpath(directory or platformdirs.user_music_path())
         if set_music_dir:
             context.set_music_dir(self.directory)
@@ -78,14 +85,14 @@ class Library(dbcore.Database):
         self._memotable: dict[tuple[str, ...], str] = {}
 
     @contextmanager
-    def music_dir_context(self):
+    def music_dir_context(self) -> Iterator[Library]:
         """Temporarily bind this library's directory to path conversion."""
         with context.music_dir(self.directory):
             yield self
 
     # Adding objects to the database.
 
-    def add(self, obj):
+    def add(self, obj: LibModel) -> int | None:
         """Add the :class:`Item` or :class:`Album` object to the library
         database.
 
@@ -95,7 +102,7 @@ class Library(dbcore.Database):
         self._memotable = {}
         return obj.id
 
-    def add_album(self, items):
+    def add_album(self, items: list[Item]) -> Album:
         """Create a new album consisting of a list of items.
 
         The items are added to the database if they don't yet have an
@@ -124,7 +131,12 @@ class Library(dbcore.Database):
 
     # Querying.
 
-    def _fetch(self, model_cls, query, sort=None):
+    def _fetch(
+        self,
+        model_cls: type[LM],
+        query: str | Sequence[str] | Query | None = None,
+        sort: Sort | None = None,
+    ) -> dbcore.Results[LM]:
         """Parse a query and fetch.
 
         If an order specification is present in the query string
@@ -151,24 +163,32 @@ class Library(dbcore.Database):
         return super()._get_results(model_cls, query, sort)
 
     @staticmethod
-    def get_default_album_sort():
+    def get_default_album_sort() -> Sort:
         """Get a :class:`Sort` object for albums from the config option."""
         return dbcore.sort_from_strings(
             Album, beets.config["sort_album"].as_str_seq()
         )
 
     @staticmethod
-    def get_default_item_sort():
+    def get_default_item_sort() -> Sort:
         """Get a :class:`Sort` object for items from the config option."""
         return dbcore.sort_from_strings(
             Item, beets.config["sort_item"].as_str_seq()
         )
 
-    def albums(self, query=None, sort=None) -> Results[Album]:
+    def albums(
+        self,
+        query: str | Sequence[str] | Query | None = None,
+        sort: Sort | None = None,
+    ) -> dbcore.Results[Album]:
         """Get :class:`Album` objects matching the query."""
         return self._fetch(Album, query, sort or self.get_default_album_sort())
 
-    def items(self, query=None, sort=None) -> Results[Item]:
+    def items(
+        self,
+        query: str | Sequence[str] | Query | None = None,
+        sort: Sort | None = None,
+    ) -> dbcore.Results[Item]:
         """Get :class:`Item` objects matching the query."""
         return self._fetch(Item, query, sort or self.get_default_item_sort())