feat: improve the type system

Author: tmgbedu

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

@@ -1,5 +1,5 @@
 import inspect
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Generic, TypeVar
 
 from fastapi_startkit.masoniteorm.expressions.expressions import (
     JoinClause,
@@ -19,9 +19,12 @@
 
 if TYPE_CHECKING:
     from fastapi_startkit.masoniteorm.connections.connection import Connection
+    from fastapi_startkit.masoniteorm.models.model import Model
 
+TModel = TypeVar("TModel", bound="Model")
 
-class QueryBuilder(EagerLoadMixin, SupportMixin):
+
+class QueryBuilder(EagerLoadMixin, SupportMixin, Generic[TModel]):
     def __init__(self, connection: "Connection", grammar, processor):
         super().__init__()
         self.connection = connection
@@ -84,26 +87,31 @@ def limit(self, limit: int) -> "QueryBuilder":
         self._limit = limit
         return self
 
-    async def find(self, primary_key: str | int, columns=None):
+    async def find(self, primary_key: str | int, columns=None) -> "TModel | None":
         return await self.where(self._model.__primary_key__, primary_key).first(columns)
 
-    async def find_or_fail(self, primary_key: str | int, columns=None):
+    async def find_or_fail(self, primary_key: str | int, columns=None) -> "TModel":
+        """Return the record matching ``primary_key``.
+
+        Raises:
+            ModelNotFoundException: If no record matches ``primary_key``.
+        """
         from fastapi_startkit.masoniteorm.exceptions import ModelNotFoundException
 
         result = await self.find(primary_key, columns)
         if result is None:
             raise ModelNotFoundException(f"{type(self._model).__name__} with primary key {primary_key!r} not found.")
         return result
 
-    async def first_or_fail(self, columns=None):
+    async def first_or_fail(self, columns=None) -> "TModel":
         from fastapi_startkit.masoniteorm.exceptions import ModelNotFoundException
 
         result = await self.first(columns)
         if result is None:
             raise ModelNotFoundException(f"{type(self._model).__name__} not found.")
         return result
 
-    async def first(self, columns=None):
+    async def first(self, columns=None) -> "TModel | None":
         if not columns:
             columns = []
 

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

@@ -1,7 +1,7 @@
 from __future__ import annotations
 import inflection
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Self
 
 from fastapi_startkit.carbon import Carbon
 from fastapi_startkit.masoniteorm.collection import Collection
@@ -92,7 +92,7 @@ def with_(cls, *eagers) -> "QueryBuilder":
         return cls.query().with_(*eagers)
 
     @classmethod
-    def where(cls, column, *args) -> "QueryBuilder":
+    def where(cls, column, *args) -> "QueryBuilder[Self]":
         return cls.query().where(column, *args)
 
     @classmethod
@@ -224,7 +224,12 @@ async def find(cls, primary_key: str | int, columns=None):
         return await cls.query().find(primary_key, columns)
 
     @classmethod
-    async def find_or_fail(cls, primary_key: str | int, columns=None):
+    async def find_or_fail(cls, primary_key: str | int, columns=None) -> Self:
+        """Fetch the record matching ``primary_key``.
+
+        Raises:
+            ModelNotFoundException: If no record matches ``primary_key``.
+        """
         return await cls.query().find_or_fail(primary_key, columns)
 
     @classmethod
@@ -280,7 +285,7 @@ def new_model_instance(self, attributes=None, exists=False):
 
         return model
 
-    def new_query(self) -> "QueryBuilder":
+    def new_query(self) -> "QueryBuilder[Self]":
         return self.db_manager.connection(self.connection).query().set_model(self)
 
     def hydrate(self, items):
@@ -310,7 +315,7 @@ def __getattr__(self, attribute):
         return self.get_attribute(attribute)
 
     @classmethod
-    def query(cls) -> "QueryBuilder":
+    def query(cls) -> "QueryBuilder[Self]":
         return cls().new_query()
 
     @classmethod