Type annotations for querysets using ClassVar

Author: davegaeddert

example/app/users/models.py

@@ -1,3 +1,7 @@
+from __future__ import annotations
+
+from typing import ClassVar
+
 from plain import models
 from plain.models import types
 from plain.passwords.types import PasswordField
@@ -8,3 +12,5 @@ class User(models.Model):
     email: str = types.EmailField()
     password: str = PasswordField()
     is_admin: bool = types.BooleanField(default=False)
+
+    query: ClassVar[models.QuerySet[User]] = models.QuerySet()

plain-api/plain/api/models.py

@@ -1,7 +1,10 @@
+from __future__ import annotations
+
 import binascii
 import os
 import uuid
 from datetime import datetime
+from typing import ClassVar
 from uuid import UUID
 
 from plain import models
@@ -26,6 +29,8 @@ class APIKey(models.Model):
 
     api_version: str = types.CharField(max_length=255, required=False)
 
+    query: ClassVar[models.QuerySet[APIKey]] = models.QuerySet()
+
     model_options = models.Options(
         constraints=[
             models.UniqueConstraint(

plain-cache/plain/cache/models.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from datetime import datetime
-from typing import Any, Self
+from typing import Any, ClassVar, Self
 
 from plain import models
 from plain.models import types
@@ -27,7 +27,7 @@ class CachedItem(models.Model):
     created_at: datetime = types.DateTimeField(auto_now_add=True)
     updated_at: datetime = types.DateTimeField(auto_now=True)
 
-    query = CachedItemQuerySet()
+    query: ClassVar[CachedItemQuerySet] = CachedItemQuerySet()
 
     model_options = models.Options(
         indexes=[

plain-flags/plain/flags/models.py

@@ -2,6 +2,7 @@
 
 import re
 from datetime import datetime
+from typing import ClassVar
 
 from plain import models
 from plain.exceptions import ValidationError
@@ -21,6 +22,8 @@ class FlagResult(models.Model):
     key: str = types.CharField(max_length=255)
     value = types.JSONField()
 
+    query: ClassVar[models.QuerySet[FlagResult]] = models.QuerySet()
+
     model_options = models.Options(
         constraints=[
             models.UniqueConstraint(
@@ -49,6 +52,8 @@ class Flag(models.Model):
     # To provide an easier way to see if a flag is still being used
     used_at: datetime | None = types.DateTimeField(required=False, allow_null=True)
 
+    query: ClassVar[models.QuerySet[Flag]] = models.QuerySet()
+
     model_options = models.Options(
         constraints=[
             models.UniqueConstraint(

plain-jobs/plain/jobs/models.py

@@ -3,7 +3,7 @@
 import datetime
 import logging
 import traceback
-from typing import TYPE_CHECKING, Any, Self
+from typing import TYPE_CHECKING, Any, ClassVar, Self
 from uuid import UUID, uuid4
 
 from opentelemetry import trace
@@ -72,6 +72,8 @@ class JobRequest(models.Model):
 
     # expires_at = models.DateTimeField(required=False, allow_null=True)
 
+    query: ClassVar[models.QuerySet[JobRequest]] = models.QuerySet()
+
     model_options = models.Options(
         ordering=["priority", "-created_at"],
         indexes=[
@@ -180,7 +182,7 @@ class JobProcess(models.Model):
         max_length=18, required=False, allow_null=True
     )
 
-    query = JobQuerySet()
+    query: ClassVar[JobQuerySet] = JobQuerySet()
 
     model_options = models.Options(
         ordering=["-created_at"],
@@ -499,7 +501,7 @@ class JobResult(models.Model):
         max_length=18, required=False, allow_null=True
     )
 
-    query = JobResultQuerySet()
+    query: ClassVar[JobResultQuerySet] = JobResultQuerySet()
 
     model_options = models.Options(
         ordering=["-created_at"],

plain-models/plain/models/README.md

@@ -7,7 +7,7 @@
 - [Querying](#querying)
 - [Migrations](#migrations)
 - [Fields](#fields)
-- [Typed field definitions](#typed-field-definitions)
+- [Typing](#typing)
 - [Validation](#validation)
 - [Indexes and constraints](#indexes-and-constraints)
 - [Custom QuerySets](#custom-querysets)
@@ -174,7 +174,7 @@ Common field types include:
 - [`URLField`](./fields/__init__.py#URLField)
 - [`UUIDField`](./fields/__init__.py#UUIDField)
 
-## Typed field definitions
+## Typing
 
 For better IDE support and type checking, use `plain.models.types` with type annotations:
 
@@ -203,6 +203,28 @@ author: Author = types.ForeignKey(Author, on_delete=models.CASCADE)
 
 All field types from the [Fields](#fields) section are available through [`types`](./types.py). Typed and untyped fields can be mixed in the same model. The database behavior is identical - typed fields only add type checking.
 
+### Typing QuerySets
+
+For better type checking of query results, you can explicitly type the `query` attribute using `ClassVar`:
+
+```python
+from __future__ import annotations
+
+from typing import ClassVar
+
+from plain import models
+from plain.models import types
+
+@models.register_model
+class User(models.Model):
+    email: str = types.EmailField()
+    is_admin: bool = types.BooleanField(default=False)
+
+    query: ClassVar[models.QuerySet[User]] = models.QuerySet()
+```
+
+With this annotation, type checkers will know that `User.query.get()` returns a `User` instance and `User.query.filter()` returns `QuerySet[User]`. The `ClassVar` annotation tells type checkers that `query` is a class-level attribute, not an instance field. This is optional - the query attribute works without the annotation, but adding it improves IDE autocomplete and type checking.
+
 ## Validation
 
 Models can be validated before saving:

plain-models/plain/models/base.py

@@ -4,7 +4,7 @@
 import warnings
 from collections.abc import Iterable, Iterator, Sequence
 from itertools import chain
-from typing import TYPE_CHECKING, Any, dataclass_transform
+from typing import TYPE_CHECKING, Any, ClassVar, dataclass_transform
 
 if TYPE_CHECKING:
     from plain.models.meta import Meta
@@ -95,7 +95,7 @@ class Model(metaclass=ModelBase):
     id = PrimaryKeyField()
 
     # Descriptors for other model behavior
-    query = QuerySet()
+    query: ClassVar[QuerySet[Model]] = QuerySet()
     model_options = Options()
     _model_meta = Meta()
     DoesNotExist = DoesNotExistDescriptor()

plain-models/plain/models/query.py

@@ -11,7 +11,7 @@
 from collections.abc import Callable, Iterator
 from functools import cached_property
 from itertools import chain, islice
-from typing import TYPE_CHECKING, Any, Generic, Self, TypeVar, overload
+from typing import TYPE_CHECKING, Any, Generic, Never, Self, TypeVar, overload
 
 import plain.runtime
 from plain.exceptions import ValidationError
@@ -338,6 +338,12 @@ def from_model(cls, model: type[T], query: sql.Query | None = None) -> Self:
         instance._deferred_filter = None
         return instance
 
+    @overload
+    def __get__(self, instance: None, owner: type[T]) -> Self: ...
+
+    @overload
+    def __get__(self, instance: Model, owner: type[T]) -> Never: ...
+
     def __get__(self, instance: Any, owner: type[T]) -> Self:
         """Descriptor protocol - return a new QuerySet bound to the model."""
         if instance is not None:

plain-oauth/plain/oauth/models.py

@@ -1,5 +1,7 @@
+from __future__ import annotations
+
 import datetime
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, ClassVar
 
 from plain import models
 from plain.auth import get_user_model
@@ -42,6 +44,8 @@ class OAuthConnection(models.Model):
         required=False, allow_null=True
     )
 
+    query: ClassVar[models.QuerySet[OAuthConnection]] = models.QuerySet()
+
     model_options = models.Options(
         constraints=[
             models.UniqueConstraint(
@@ -71,13 +75,13 @@ def refresh_access_token(self) -> None:
         self.set_token_fields(refreshed_oauth_token)
         self.save()
 
-    def set_token_fields(self, oauth_token: "OAuthToken") -> None:
+    def set_token_fields(self, oauth_token: OAuthToken) -> None:
         self.access_token = oauth_token.access_token
         self.refresh_token = oauth_token.refresh_token
         self.access_token_expires_at = oauth_token.access_token_expires_at
         self.refresh_token_expires_at = oauth_token.refresh_token_expires_at
 
-    def set_user_fields(self, oauth_user: "OAuthUser") -> None:
+    def set_user_fields(self, oauth_user: OAuthUser) -> None:
         self.provider_user_id = oauth_user.provider_id
 
     def access_token_expired(self) -> bool:
@@ -94,8 +98,8 @@ def refresh_token_expired(self) -> bool:
 
     @classmethod
     def get_or_create_user(
-        cls, *, provider_key: str, oauth_token: "OAuthToken", oauth_user: "OAuthUser"
-    ) -> "OAuthConnection":
+        cls, *, provider_key: str, oauth_token: OAuthToken, oauth_user: OAuthUser
+    ) -> OAuthConnection:
         try:
             connection = cls.query.get(
                 provider_key=provider_key,
@@ -129,9 +133,9 @@ def connect(
         *,
         user: Any,
         provider_key: str,
-        oauth_token: "OAuthToken",
-        oauth_user: "OAuthUser",
-    ) -> "OAuthConnection":
+        oauth_token: OAuthToken,
+        oauth_user: OAuthUser,
+    ) -> OAuthConnection:
         """
         Connect will either create a new connection or update an existing connection
         """

plain-observer/plain/observer/models.py

@@ -6,7 +6,7 @@
 from collections.abc import Iterable, Mapping, Sequence
 from datetime import UTC, datetime
 from functools import cached_property
-from typing import TYPE_CHECKING, Any, cast
+from typing import TYPE_CHECKING, Any, ClassVar, cast
 
 import sqlparse
 from opentelemetry.sdk.trace import ReadableSpan
@@ -66,6 +66,8 @@ class Trace(models.Model):
         spans: BaseRelatedManager
         logs: BaseRelatedManager
 
+    query: ClassVar[models.QuerySet[Trace]] = models.QuerySet()
+
     model_options = models.Options(
         ordering=["-start_time"],
         constraints=[
@@ -335,7 +337,7 @@ class Span(models.Model):
     status: str = types.CharField(max_length=50, default="", required=False)
     span_data: dict = types.JSONField(default=dict, required=False)
 
-    query = SpanQuerySet()
+    query: ClassVar[SpanQuerySet] = SpanQuerySet()
 
     model_options = models.Options(
         ordering=["-start_time"],
@@ -523,6 +525,8 @@ class Log(models.Model):
     level: str = types.CharField(max_length=20)
     message: str = types.TextField()
 
+    query: ClassVar[models.QuerySet[Log]] = models.QuerySet()
+
     model_options = models.Options(
         ordering=["timestamp"],
         indexes=[