Skip to content

Commit 46c3748

Browse files
tmgbeduclaude
andcommitted
feat(orm): type QueryBuilder.get() as Collection[TModel]
Make the ORM Collection generic (Collection[T]) with TYPE_CHECKING-only element-access stubs (first/__iter__/__getitem__), and annotate QueryBuilder.get()/get_models() to return Collection[TModel]. So `User.where(...).get()` now types as Collection[User] instead of Any, and element access (iteration, indexing, first()) yields User. - load() gets one `# type: ignore` (runs only on model collections, but the ORM Collection is legitimately generic over non-model values too). - chunk_by_id uses results[-1] (equivalent; non-emptiness already guaranteed) so the last row is cleanly typed. Stubs are typing-only; runtime behaviour is unchanged. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1 parent 1fe875d commit 46c3748

2 files changed

Lines changed: 19 additions & 4 deletions

File tree

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

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,19 @@
1+
from typing import TYPE_CHECKING, Any, Generator, Generic, TypeVar
2+
13
from fastapi_startkit.support.collection import Collection as BaseCollection
24

5+
T = TypeVar("T")
6+
7+
8+
class Collection(BaseCollection, Generic[T]):
9+
if TYPE_CHECKING:
10+
# Typing-only element-access overrides so a Collection[User] yields
11+
# User (not Any) on iteration, indexing, and first(). Runtime behaviour
12+
# is supplied unchanged by the base class.
13+
def first(self, callback=None) -> "T | None": ...
14+
def __iter__(self) -> "Generator[T, Any, None]": ...
15+
def __getitem__(self, item) -> "T": ...
316

4-
class Collection(BaseCollection):
517
def with_relationship_autoloading(self):
618
pass
719

@@ -27,6 +39,7 @@ async def load(self, *relations):
2739
if isinstance(result_set, Collection):
2840
relationship.register_related(relation, model, map_related)
2941
else:
30-
model.add_relation({relation: map_related or None})
42+
# load() only runs on model collections; T is generic.
43+
model.add_relation({relation: map_related or None}) # type: ignore
3144

3245
return self

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

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
from fastapi_startkit.masoniteorm.query.support import SupportMixin
1919

2020
if TYPE_CHECKING:
21+
from fastapi_startkit.masoniteorm.collection import Collection
2122
from fastapi_startkit.masoniteorm.connections.connection import Connection
2223
from fastapi_startkit.masoniteorm.models.model import Model
2324

@@ -158,7 +159,7 @@ async def first(self, columns=None) -> "TModel | None":
158159
results = await self.select(columns).limit(1).get()
159160
return results.first()
160161

161-
async def get(self, columns=None):
162+
async def get(self, columns=None) -> "Collection[TModel]":
162163
# TODO: apply scopes
163164
if not columns:
164165
columns = []
@@ -458,7 +459,8 @@ async def chunk_by_id(self, count: int, column: str = None, alias: str = None, d
458459

459460
yield results
460461

461-
last_id = results.last().get_attributes().get(alias)
462+
# count_results != 0 above guarantees a last row.
463+
last_id = results[-1].get_attributes().get(alias)
462464
if last_id is None:
463465
raise RuntimeError(
464466
f"The chunk_by_id operation was aborted because the [{alias}] "

0 commit comments

Comments
 (0)