diff --git a/fastapi_startkit/src/fastapi_startkit/masoniteorm/collection/Collection.py b/fastapi_startkit/src/fastapi_startkit/masoniteorm/collection/Collection.py index 8e4c492d..2882b6b0 100644 --- a/fastapi_startkit/src/fastapi_startkit/masoniteorm/collection/Collection.py +++ b/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 diff --git a/fastapi_startkit/src/fastapi_startkit/masoniteorm/models/builder.py b/fastapi_startkit/src/fastapi_startkit/masoniteorm/models/builder.py index 45d45c59..e822f0af 100644 --- a/fastapi_startkit/src/fastapi_startkit/masoniteorm/models/builder.py +++ b/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}] " diff --git a/fastapi_startkit/tests/masoniteorm/collection/test_collection.py b/fastapi_startkit/tests/masoniteorm/collection/test_collection.py index 873b4837..5bbe1d7b 100644 --- a/fastapi_startkit/tests/masoniteorm/collection/test_collection.py +++ b/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())