Add utilities for working with id lists

src/utils/db.py

@@ -1,6 +1,7 @@
-from typing import final
+from collections.abc import Sequence
+from typing import Literal, final, overload
 
-from beanie import Document, Link
+from beanie import Document, Link, PydanticObjectId
 
 
 @final
@@ -25,3 +26,62 @@ class UnfetchedLinkError[T](ValueError):
     def __init__(self, link: Link[T]):
         self.link = link
         super().__init__(f"Link {link} was not fetched")
+
+
+@final
+class IdNotFoundError[T: Document](ValueError):
+    """Raised when an object with given id was not found."""
+
+    def __init__(self, item_id: PydanticObjectId, typ: type[T]) -> None:
+        self.item_id = item_id
+        self.typ = typ
+        super().__init__(f"Item with id {item_id} of type {typ.__name__} not found")
+
+
+def get_id_list[T: Document](items: Sequence[T | Link[T]]) -> list[PydanticObjectId]:
+    """Extracts the ids from a list of db items."""
+    id_list: list[PydanticObjectId] = []
+    for item in items:
+        if isinstance(item, Link):
+            raise UnfetchedLinkError(item)
+
+        if item.id is None:
+            raise MissingIdError(item)
+
+        id_list.append(item.id)
+
+    return id_list
+
+
+@overload
+async def from_id_list[T: Document](
+    id_list: Sequence[PydanticObjectId],
+    typ: type[T],
+    raise_on_missing: Literal[True] = True,
+) -> list[T]: ...
+
+
+@overload
+async def from_id_list[T: Document](
+    id_list: Sequence[PydanticObjectId],
+    typ: type[T],
+    raise_on_missing: Literal[False],
+) -> list[T | None]: ...
+
+
+async def from_id_list[T: Document](
+    id_list: Sequence[PydanticObjectId],
+    typ: type[T],
+    raise_on_missing: bool = True,
+) -> list[T] | list[T | None]:
+    """Constructs a list of db objects from a list of ids."""
+    items: list[T | None] = []
+    for item_id in id_list:
+        item = await typ.get(item_id)
+        if item is None:
+            if raise_on_missing:
+                raise IdNotFoundError(item_id, typ)
+            items.append(None)
+            continue
+        items.append(item)
+    return items