Type AsyncQueryBuilder.where() and .get() so basedpyright resolves concrete types

fastapi_startkit/pyrightconfig.typing-tests.json

@@ -0,0 +1,7 @@
+{
+  "include": [
+    "tests/masoniteorm/sqlite/models/test_model_query_typing.py"
+  ],
+  "typeCheckingMode": "standard",
+  "pythonVersion": "3.12"
+}

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

@@ -1,4 +1,5 @@
-from typing import TYPE_CHECKING, Any, Generator, Generic, TypeVar
+from collections.abc import Callable, Generator
+from typing import TYPE_CHECKING, Any, Generic, TypeVar
 
 from fastapi_startkit.support.collection import Collection as BaseCollection
 
@@ -10,7 +11,7 @@ class Collection(BaseCollection, Generic[T]):
         # 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 first(self, callback: "Callable[[T], bool] | None" = None) -> "T | None": ...
         def __iter__(self) -> "Generator[T, Any, None]": ...
         def __getitem__(self, item) -> "T": ...
 

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

@@ -1,5 +1,6 @@
 import inspect
-from typing import TYPE_CHECKING, Any, Generic, TypeVar
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any, Generic, Self, TypeVar, overload
 
 from fastapi_startkit.masoniteorm.expressions.expressions import (
     JoinClause,
@@ -24,6 +25,10 @@
 
 TModel = TypeVar("TModel", bound="Model")
 
+# where(lambda q: q.where(...)) — the callable receives a nested builder and
+# returns it, which the parent renders as a parenthesised subgroup.
+type WhereGroup[M: "Model"] = Callable[["QueryBuilder[M]"], "QueryBuilder[M]"]
+
 
 class QueryBuilder(EagerLoadMixin, SupportMixin, Generic[TModel]):
     operators = [
@@ -159,7 +164,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) -> "Collection[TModel]":
+    async def get(self, columns: "list[str] | str | None" = None) -> "Collection[TModel]":
         # TODO: apply scopes
         if not columns:
             columns = []
@@ -486,11 +491,27 @@ def invalid_operator(self, operator):
         """Determine whether an operator is not supported by the builder."""
         return not isinstance(operator, str) or operator.lower() not in self.operators
 
-    def where(self, column, *args):
+    @overload
+    def where(self, column: str, /) -> "Self": ...
+
+    @overload
+    def where(self, column: str, value: Any, /) -> "Self": ...
+
+    @overload
+    def where(self, column: str, operator: str, value: Any, /) -> "Self": ...
+
+    @overload
+    def where(self, column: dict[str, Any], /) -> "Self": ...
+
+    @overload
+    def where(self, column: "WhereGroup[TModel]", /) -> "Self": ...
+
+    def where(self, column: "str | dict[str, Any] | WhereGroup[TModel]", *args: Any) -> "Self":
         """Specifies a where expression.
 
         Arguments:
-            column {string} -- The name of the column to search
+            column {string | dict | callable} -- The column to search, a dict of
+                column/value pairs, or a callable receiving a nested builder.
 
         Keyword Arguments:
             args {List} -- The operator and the value of the column to search. (default: {None})

fastapi_startkit/src/fastapi_startkit/masoniteorm/models/model.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Self
+from typing import TYPE_CHECKING, Any, Self, overload
 
 import inflection
 import pendulum
@@ -15,7 +15,7 @@
 from fastapi_startkit.masoniteorm.observers import ObservesEvents
 
 if TYPE_CHECKING:
-    from fastapi_startkit.masoniteorm.models.builder import QueryBuilder
+    from fastapi_startkit.masoniteorm.models.builder import QueryBuilder, WhereGroup
 
 
 class Model(Attribute, Relationship, ObservesEvents):
@@ -93,8 +93,28 @@ def get_related(self, key: str):
     def with_(cls, *eagers) -> "QueryBuilder":
         return cls.query().with_(*eagers)
 
+    @overload
     @classmethod
-    def where(cls, column, *args) -> "QueryBuilder[Self]":
+    def where(cls, column: str, /) -> QueryBuilder[Self]: ...
+
+    @overload
+    @classmethod
+    def where(cls, column: str, value: Any, /) -> QueryBuilder[Self]: ...
+
+    @overload
+    @classmethod
+    def where(cls, column: str, operator: str, value: Any, /) -> QueryBuilder[Self]: ...
+
+    @overload
+    @classmethod
+    def where(cls, column: dict[str, Any], /) -> QueryBuilder[Self]: ...
+
+    @overload
+    @classmethod
+    def where(cls, column: WhereGroup[Self], /) -> QueryBuilder[Self]: ...
+
+    @classmethod
+    def where(cls, column: str | dict[str, Any] | WhereGroup[Self], *args: Any) -> QueryBuilder[Self]:
         return cls.query().where(column, *args)
 
     @classmethod
@@ -243,7 +263,7 @@ async def first(cls, columns=None):
         return await cls.query().first(columns)
 
     @classmethod
-    async def get(cls):
+    async def get(cls) -> Collection[Self]:
         return await cls.query().get()
 
     @classmethod

fastapi_startkit/tests/masoniteorm/sqlite/models/test_model_query_typing.py

@@ -0,0 +1,86 @@
+# pyright: reportUnknownMemberType=error, reportUnknownVariableType=error, reportUnknownArgumentType=error
+"""Static typing guarantees of where() and get() on models and the query builder.
+
+The type checker is the real assertion here:
+
+    cd fastapi_startkit && uv run pyright -p pyrightconfig.typing-tests.json
+
+(the repo pyright config excludes tests/, so this file ships its own config).
+
+`assert_type` is a no-op at runtime, so every static assertion is paired with a
+runtime check that the value really is what the annotation promises. The
+`reportUnknown*` rules above are enabled per-file: an unannotated parameter on
+where()/get() makes this module fail to type-check even though the assert_type
+calls themselves would still pass.
+"""
+
+from typing import assert_type
+
+from fastapi_startkit.masoniteorm.collection import Collection
+from fastapi_startkit.masoniteorm.models.builder import QueryBuilder
+
+from ...fixtures.model import User
+from ..test_case import TestCase
+
+
+class TestWhereTyping(TestCase):
+    async def test_value_form_returns_builder_of_the_model(self):
+        builder = User.where("name", "Joe")
+
+        assert_type(builder, QueryBuilder[User])
+        assert isinstance(builder, QueryBuilder)
+
+    async def test_operator_form_returns_builder_of_the_model(self):
+        builder = User.where("name", "!=", "Joe")
+
+        assert_type(builder, QueryBuilder[User])
+        assert isinstance(builder, QueryBuilder)
+
+    async def test_dict_form_returns_builder_of_the_model(self):
+        builder = User.where({"name": "Joe", "is_admin": True})
+
+        assert_type(builder, QueryBuilder[User])
+        assert isinstance(builder, QueryBuilder)
+
+    async def test_callable_form_returns_builder_of_the_model(self):
+        builder = User.query().where(lambda q: q.where("name", "Joe").where("is_admin", True))
+
+        assert_type(builder, QueryBuilder[User])
+        assert isinstance(builder, QueryBuilder)
+
+    async def test_model_type_survives_a_multi_step_chain(self):
+        builder = User.where("is_admin", True).where("name", "!=", "Jane").where("name", "like", "%o%")
+
+        assert_type(builder, QueryBuilder[User])
+        assert isinstance(builder, QueryBuilder)
+
+
+class TestGetTyping(TestCase):
+    async def test_awaited_get_is_a_collection_of_the_model(self):
+        users = await User.where("name", "Joe").get()
+
+        assert_type(users, Collection[User])
+        assert isinstance(users, Collection)
+        assert len(users) == 1
+
+    async def test_collection_elements_are_the_model(self):
+        users = await User.where("name", "Joe").get()
+        user = users.first()
+
+        assert_type(user, User | None)
+        assert isinstance(user, User)
+        assert user.name == "Joe"
+
+    async def test_get_with_explicit_columns(self):
+        users = await User.where("name", "Joe").get(["name"])
+
+        assert_type(users, Collection[User])
+        user = users.first()
+        assert user is not None
+        assert user.name == "Joe"
+
+    async def test_model_get_is_a_collection_of_the_model(self):
+        users = await User.get()
+
+        assert_type(users, Collection[User])
+        assert len(users) == 2