feat(orm): type QueryBuilder.get() as Collection[TModel]

fastapi_startkit/src/fastapi_startkit/masoniteorm/collection/Collection.py

@@ -1,7 +1,19 @@
+from typing import TYPE_CHECKING, Any, Generator, Generic, TypeVar
+
 from fastapi_startkit.support.collection import Collection as BaseCollection
 
+T = TypeVar("T")
+
+
+class Collection(BaseCollection, Generic[T]):
+    if TYPE_CHECKING:
+        # Typing-only element-access overrides so a Collection[User] yields
+        # User (not Any) on iteration, indexing, and first(). Runtime behaviour
+        # is supplied unchanged by the base class.
+        def first(self, callback=None) -> "T | None": ...
+        def __iter__(self) -> "Generator[T, Any, None]": ...
+        def __getitem__(self, item) -> "T": ...
 
-class Collection(BaseCollection):
     def with_relationship_autoloading(self):
         pass
 
@@ -27,6 +39,7 @@ async def load(self, *relations):
                     if isinstance(result_set, Collection):
                         relationship.register_related(relation, model, map_related)
                     else:
-                        model.add_relation({relation: map_related or None})
+                        # load() only runs on model collections; T is generic.
+                        model.add_relation({relation: map_related or None})  # type: ignore
 
         return self

fastapi_startkit/src/fastapi_startkit/masoniteorm/models/builder.py

@@ -18,6 +18,7 @@
 from fastapi_startkit.masoniteorm.query.support import SupportMixin
 
 if TYPE_CHECKING:
+    from fastapi_startkit.masoniteorm.collection import Collection
     from fastapi_startkit.masoniteorm.connections.connection import Connection
     from fastapi_startkit.masoniteorm.models.model import Model
 
@@ -158,7 +159,7 @@ async def first(self, columns=None) -> "TModel | None":
         results = await self.select(columns).limit(1).get()
         return results.first()
 
-    async def get(self, columns=None):
+    async def get(self, columns=None) -> "Collection[TModel]":
         # TODO: apply scopes
         if not columns:
             columns = []
@@ -458,7 +459,8 @@ async def chunk_by_id(self, count: int, column: str = None, alias: str = None, d
 
             yield results
 
-            last_id = results.last().get_attributes().get(alias)
+            # count_results != 0 above guarantees a last row.
+            last_id = results[-1].get_attributes().get(alias)
             if last_id is None:
                 raise RuntimeError(
                     f"The chunk_by_id operation was aborted because the [{alias}] "

fastapi_startkit/tests/masoniteorm/collection/test_collection.py

@@ -65,6 +65,54 @@ async def test_load_on_empty_collection_returns_self(self):
 
         self.assertIs(result, users)
 
+    async def test_load_registers_non_collection_result_via_add_relation(self):
+        # Defensive branch in load(): when a relationship's get_related returns
+        # a truthy value that is NOT a Collection, load() attaches it to every
+        # model with add_relation instead of the register_related batch path.
+        # No built-in relationship hits this (they all return Collections for a
+        # collection input), so it is exercised with a stub relationship.
+        class StubRelationship:
+            async def get_related(self, query, relation, eagers=None, callback=None):
+                return {"payload": 1}  # truthy, not a Collection
+
+            def map_related(self, related_result):
+                return related_result
+
+            def register_related(self, key, model, collection):
+                raise AssertionError("register_related must not run for a non-Collection result")
+
+        class Widget(Model):
+            gadget = StubRelationship()
+
+        items = Collection([Widget(), Widget()])
+
+        result = await items.load("gadget")
+
+        self.assertIs(result, items)
+        for widget in items:
+            self.assertEqual(widget._relationships["gadget"], {"payload": 1})
+
+    async def test_load_registers_none_when_mapped_result_is_falsy(self):
+        # Same branch, but a falsy mapped result is stored as None.
+        class StubRelationship:
+            async def get_related(self, query, relation, eagers=None, callback=None):
+                return "truthy-raw-result"  # not a Collection, so else branch
+
+            def map_related(self, related_result):
+                return {}  # falsy → `map_related or None` stores None
+
+            def register_related(self, key, model, collection):
+                raise AssertionError("register_related must not run here")
+
+        class Gizmo(Model):
+            part = StubRelationship()
+
+        items = Collection([Gizmo()])
+
+        await items.load("part")
+
+        self.assertIsNone(items.first()._relationships["part"])
+
     def test_with_relationship_autoloading_is_noop(self):
         self.assertIsNone(Collection([]).with_relationship_autoloading())