diff --git a/plain-postgres/plain/postgres/README.md b/plain-postgres/plain/postgres/README.md index b69ec4f8e8..da2fbd7a95 100644 --- a/plain-postgres/plain/postgres/README.md +++ b/plain-postgres/plain/postgres/README.md @@ -258,6 +258,77 @@ A traversed field _is_ the related field, carrying the relation path as its name [Encrypted fields](#encrypted-fields) reject value comparisons because their ciphertext is non-deterministic — only `is_null()` is available, and any other condition method (`equals`, `is_in`, …) raises `TypeError`. +### Selecting columns with select() + +`select()` pulls back specific columns as typed rows instead of model instances. You pass typed field references, and a type checker knows the exact shape of each row: + +```python +from plain.postgres import Field, types + + +@postgres.register_model +class User(postgres.Model): + email: Field[str] = types.EmailField() + age: Field[int | None] = types.IntegerField(allow_null=True, default=None) + + +# list-like of tuple[str, int | None], precisely typed +rows = User.query.where(User.age.gte(18)).select(User.email, User.age) +for email, age in rows: + ... +``` + +There are three modes: + +- **Tuples** (default) — one tuple per row, typed per column: `select(User.email, User.age)` yields `tuple[str, int | None]`. +- **Flat scalars** — a single column unwrapped, with `flat=True`: `select(User.email, flat=True)` yields `str`. `flat=True` accepts exactly one column. +- **Dataclasses** — map each column onto a dataclass with `result_type=`: `select(User.email, User.age, result_type=UserStats)` yields `UserStats`. Columns map to dataclass fields **positionally**, so the selection order must match the dataclass field order, and each selected field's name must match the dataclass field at the same position. + +```python +from dataclasses import dataclass + + +@dataclass +class UserStats: + email: str + age: int | None + + +stats = User.query.select(User.email, User.age, result_type=UserStats) +``` + +The return value is a [`RowQuerySet`](./query.py#RowQuerySet) — that is the name to reach for when you need to annotate one: + +```python +from plain.postgres import RowQuerySet + + +def adults() -> RowQuerySet[tuple[str, int | None]]: + return User.query.where(User.age.gte(18)).select(User.email, User.age) +``` + +You can select expression columns too — `select(User.id, Sum("amount"))`, or an `F()` — but an expression column types as `Any` (its output type isn't tracked yet). The fields around it stay precise, so `select(User.id, Sum("amount"))` types as `tuple[int, Any]`. + +Per-column typing runs to **ten columns**. An eleventh is still selected and still returns rows, but the row type degrades to `tuple[Any, ...]` — reach for `result_type=` when a row is that wide. + +`select()` goes last in a chain: `annotate()` must come before it, because an annotation appends a column and would change the row shape out from under the type `select()` declared. `annotate()` after `select()` raises `TypeError` saying so. `prefetch_related()` is refused in both orders — a prefetch attaches related objects to a model instance's attributes, and a row has nowhere to put them; select the columns you need from the related model instead. + +Re-selecting replaces the **column list**, not the joins: an expression that reached through a relation (`select(Upper("tags__name"))`) leaves its join in place, so a later `select(Widget.name)` still returns one row per joined row — and `count()`/`exists()` count those. This is `annotate(...)` followed by `values_list(...)` behaving as it always has; trimming joins no queryset needs any more is out of scope here. + +`distinct()` with an `order_by()` on a column you didn't select returns duplicates: the ordering column has to go into the `SELECT` list for Postgres to sort by it, so `SELECT DISTINCT` deduplicates on that column too. Order by something you selected, or drop the ordering. This is `values_list()`'s behavior as well, not new to `select()`. + +Columns annotated `Field[Any]` are rejected by `select()`, because `Any` satisfies the model-valued `__get__` overload and class access resolves as `type[Any]` rather than a field. That's the same reason the field-annotation guidance says never to annotate a field `Field[Any]` — use the concrete type, or `Field[object]` when the column really does hold arbitrary JSON, which `select()` types as `object`. + +**`select()` returns rows, not partial model instances.** This is deliberate: a model instance with only some columns loaded is a type-level lie — the type checker thinks every field is present, so touching an unselected column looks fine but fails or fires a hidden query at runtime. Honest tuples/dataclasses keep the types truthful. As a result, iteration, `first()`, `get()`, `iterator()`, and slicing all return rows, and anything that would read or write model rows, or change the selected columns — `update()`, `delete()`, `get_or_create()`, `values()`, `values_list()`, `annotate()`, `prefetch_related()` — raises `TypeError`. `update()` and `delete()` refuse a queryset in row mode however it got there, `values()` and `values_list()` included. + +`select()` takes typed references only — a bare string like `select("email")` raises `TypeError` (use `User.email`). + +**A column belongs to the model whose field built it**, the same as [a condition does](#querying-with-typed-conditions): `Order.query.select(User.email)` raises `TypeError` naming both models. A type checker can't catch it — `Field[str]` is `Field[str]` whichever model declared it — and without the check the name `"email"` just resolves against `Order`, silently the wrong column when both models have one. Expressions are unaffected: `F("email")` and `Upper("email")` take a string resolved against whatever query they land in, like `filter()`'s kwargs. + +**Relations are not selectable yet.** `select(Post.author)` (the relation) and `select(Post.author.city)` (a column through it) both raise `TypeError`, and so does `select(Post.author.id)` — the foreign key column itself. The reason is nullability: a column reached through a relation arrives over a join, so a nullable relation yields `None` where the traversed field's type says it can't. Until `select()` can express that, `values_list("author__id", flat=True)` is the spelling, and the error message names it. + +**`select()` hands back a plain `RowQuerySet`, not your custom QuerySet subclass.** Chain your own methods before `select()`, not after — `User.query.active().select(...)` works, `User.query.select(...).active()` raises `AttributeError`. + ### Custom QuerySets You can customize [`QuerySet`](./query.py#QuerySet) classes to provide specialized query methods. Define a custom QuerySet and assign it to your model's `query` attribute as a `ClassVar` (so it isn't treated as a constructor field): diff --git a/plain-postgres/plain/postgres/__init__.py b/plain-postgres/plain/postgres/__init__.py index 05366993c4..86fb5a3f41 100644 --- a/plain-postgres/plain/postgres/__init__.py +++ b/plain-postgres/plain/postgres/__init__.py @@ -49,7 +49,7 @@ ) from .indexes import Index from .options import Options -from .query import QuerySet +from .query import QuerySet, RowQuerySet from .query_utils import Q from . import types @@ -98,6 +98,7 @@ "RandomStringField", "ReverseForeignKey", "ReverseManyToMany", + "RowQuerySet", "SmallIntegerField", "TextChoices", "TextField", diff --git a/plain-postgres/plain/postgres/expressions.py b/plain-postgres/plain/postgres/expressions.py index 03ede3d301..ee13c068d6 100644 --- a/plain-postgres/plain/postgres/expressions.py +++ b/plain-postgres/plain/postgres/expressions.py @@ -33,6 +33,7 @@ FullResultSet, ) from plain.postgres.query_utils import Q +from plain.postgres.selectable import Selectable from plain.utils.deconstruct import deconstructible from plain.utils.hashable import make_hashable @@ -220,7 +221,7 @@ def __invert__(self) -> NegatedExpression: return NegatedExpression(self) -class BaseExpression: +class BaseExpression(Selectable[Any]): """Base class for all query expressions.""" empty_result_set_value = NotImplemented @@ -809,7 +810,7 @@ def as_sql( @deconstructible(path="plain.postgres.F") -class F(Combinable): +class F(Combinable, Selectable[Any]): """An object capable of resolving references to existing query objects.""" def __init__(self, name: str): diff --git a/plain-postgres/plain/postgres/fields/__init__.py b/plain-postgres/plain/postgres/fields/__init__.py index 3cbc6682fb..15b74dc7be 100644 --- a/plain-postgres/plain/postgres/fields/__init__.py +++ b/plain-postgres/plain/postgres/fields/__init__.py @@ -3,7 +3,6 @@ from .base import DATABASE_DEFAULT as DATABASE_DEFAULT from .base import ( NOT_PROVIDED, - Empty, Field, ) from .base import ChoicesField as ChoicesField @@ -34,7 +33,6 @@ "DecimalField", "DurationField", "EmailField", - "Empty", "Field", "FloatField", "GenericIPAddressField", diff --git a/plain-postgres/plain/postgres/fields/base.py b/plain-postgres/plain/postgres/fields/base.py index 76141d68cd..40a6a9174f 100644 --- a/plain-postgres/plain/postgres/fields/base.py +++ b/plain-postgres/plain/postgres/fields/base.py @@ -17,6 +17,7 @@ from plain.postgres.dialect import quote_name from plain.postgres.enums import ChoicesMeta from plain.postgres.query_utils import Q, RegisterLookupMixin +from plain.postgres.selectable import Selectable from plain.preflight import PreflightResult from plain.utils.datastructures import DictWrapper from plain.utils.functional import Promise @@ -34,10 +35,6 @@ from plain.postgres.sql.compiler import SQLCompiler -class Empty: - pass - - class NOT_PROVIDED: pass @@ -91,17 +88,20 @@ def _load_field( # except for ForeignKeys, where the "_id" suffix is appended. -def _empty(of_cls: type) -> Empty: - new = Empty() - new.__class__ = of_cls - return new +def _empty(of_cls: type) -> Any: + """Build an initialized-but-unpopulated instance of `of_cls`. + + Module-level (not a lambda or a method) because `__reduce__` names it as + the pickle reconstructor. + """ + return object.__new__(of_cls) # Ordering conditions: the ones a None operand is meaningless for. _ORDERING_SUFFIXES = frozenset({"gt", "gte", "lt", "lte"}) -class Field[T](RegisterLookupMixin): +class Field[T](Selectable[T], RegisterLookupMixin): """Base class for all field types""" # SQL type for this field (e.g. "text", "integer", "boolean"). @@ -434,12 +434,11 @@ def __deepcopy__(self, memodict: dict[int, Any]) -> Self: return obj def __copy__(self) -> Self: - # We need to avoid hitting __reduce__, so define this - # slightly weird copy construct. - obj = Empty() - obj.__class__ = self.__class__ + # Build the instance directly rather than calling the constructor, + # which would hit __reduce__. + obj = object.__new__(self.__class__) obj.__dict__ = self.__dict__.copy() - return cast(Self, obj) + return obj def __reduce__( self, diff --git a/plain-postgres/plain/postgres/query.py b/plain-postgres/plain/postgres/query.py index 6486648731..313c7242d0 100644 --- a/plain-postgres/plain/postgres/query.py +++ b/plain-postgres/plain/postgres/query.py @@ -5,7 +5,9 @@ from __future__ import annotations import copy +import dataclasses import datetime +import inspect import json import operator import warnings @@ -13,7 +15,7 @@ from decimal import Decimal from functools import cached_property from itertools import islice -from typing import TYPE_CHECKING, Any, Never, Self, cast, overload +from typing import TYPE_CHECKING, Any, Literal, Never, Self, cast, overload import plain.runtime import psycopg @@ -46,6 +48,7 @@ from plain.postgres.fields.json import JSONField from plain.postgres.functions import Cast from plain.postgres.query_utils import Q, condition_origins_of +from plain.postgres.selectable import Selectable from plain.postgres.sql import ( AND, CURSOR, @@ -61,9 +64,10 @@ from plain.utils.functional import partition # Re-exports for public API -__all__ = ["F", "Prefetch", "Q", "QuerySet", "RawQuerySet"] +__all__ = ["F", "Prefetch", "Q", "QuerySet", "RawQuerySet", "RowQuerySet"] if TYPE_CHECKING: + from _typeshed import DataclassInstance from plain.postgres import Model @@ -363,6 +367,41 @@ def __iter__(self) -> Iterator[Any]: yield row[0] +class SelectDataclassIterable(BaseIterable): + """ + Iterable returned by QuerySet.select(result_type=...) that builds one + dataclass instance per row. Columns map to dataclass fields positionally, + so the tuple rows from ValuesListIterable are zipped onto the dataclass + field names in order. + """ + + def __iter__(self) -> Iterator[Any]: + queryset = self.queryset + result_type = queryset._select_result_type + assert result_type is not None + tuple_rows = ValuesListIterable(queryset, chunked_fetch=self.chunked_fetch) + + # Work out how to call the constructor once, not once per row. + # Parameter kind decides how each value has to be passed: a + # keyword-only one can't be filled positionally, and a positional-only + # one can't be filled by name. A signature always orders positionals + # before keyword-onlys, so the row splits at a single point. + parameters = _result_type_parameters(result_type) + keyword_names = tuple( + p.name for p in parameters if p.kind is inspect.Parameter.KEYWORD_ONLY + ) + if not keyword_names: + for row in tuple_rows: + yield result_type(*row) + return + + split = len(parameters) - len(keyword_names) + for row in tuple_rows: + yield result_type( + *row[:split], **dict(zip(keyword_names, row[split:], strict=True)) + ) + + class QuerySet[T: "Model"]: """ Represent a lazy database lookup for a set of objects. @@ -396,6 +435,8 @@ class Task(Model): _known_related_objects: dict[Any, dict[Any, Any]] _iterable_class: type[BaseIterable] _fields: tuple[str, ...] | None + # Set by select(result_type=...); drives SelectDataclassIterable. + _select_result_type: type[DataclassInstance] | None _defer_next_filter: bool _deferred_filter: tuple[bool, tuple[Any, ...], dict[str, Any]] | None # The columns to RETURN from the next update()/delete(), or None for a @@ -422,6 +463,7 @@ def from_model(cls, model: type[T], query: Query | None = None) -> Self: instance._known_related_objects = {} instance._iterable_class = ModelIterable instance._fields = None + instance._select_result_type = None instance._defer_next_filter = False instance._deferred_filter = None instance._returning_fields = None @@ -627,14 +669,16 @@ def __or__(self, other: QuerySet[T]) -> Any: self if self.sql_query.can_filter() else self.model._model_meta.base_queryset.filter( - id__in=self.values("id") + id__in=self._values("id") ) ) combined = cast("Self", query._chain()) combined._merge_known_related_objects(other) if not other.sql_query.can_filter(): other = other.model._model_meta.base_queryset.filter( - id__in=other.values("id") + # `_values`, not `values()`: these are subqueries, and the + # public method is refused on a row-mode queryset. + id__in=other._values("id") ) combined.sql_query.combine(other.sql_query, OR) combined._returning_fields, combined._returning_instances = returning @@ -1737,7 +1781,12 @@ def update(self, **kwargs: Any) -> int: if self.sql_query.is_sliced: raise TypeError("Cannot update a query once a slice has been taken.") if self._fields is not None: - raise TypeError("Cannot call update() after .values() or .values_list()") + # Same guard delete() carries: once the queryset is in row mode + # (values(), values_list(), select()) it no longer describes the + # model rows a write would touch. + raise TypeError( + "Cannot call update() after .values(), .values_list() or .select()" + ) self._reject_related_lock_targets("update") query = self.sql_query.chain(UpdateQuery) query.add_update_values(kwargs) @@ -1835,7 +1884,10 @@ def raw( def _values(self, *fields: str, **expressions: Any) -> QuerySet[Any]: clone = self._chain() if expressions: - clone = clone.annotate(**expressions) + # The internal mechanism, not the public method: select() aliases + # its expression columns through here, and RowQuerySet refuses + # annotate(). + clone = clone._annotate(**expressions) clone._fields = fields clone.sql_query.set_values(list(fields)) return clone @@ -1847,13 +1899,35 @@ def values(self, *fields: str, **expressions: Any) -> QuerySet[Any]: return clone def values_list(self, *fields: str, flat: bool = False) -> QuerySet[Any]: + return self._values_list(fields, flat=flat) + + def _values_list( + self, fields: tuple[str | ResolvableExpression, ...], *, flat: bool + ) -> QuerySet[Any]: if flat and len(fields) > 1: raise TypeError( "'flat' is not valid when values_list is called with more than one " "field." ) - field_names = {f for f in fields if not isinstance(f, ResolvableExpression)} + # Names an internal alias must not collide with. The newly selected + # columns are the obvious ones, but the counter also has to clear + # everything already on the query: + # + # * `self.sql_query.annotations` -- a user's own `annotate(upper1=...)` + # would otherwise be silently overwritten, quietly changing what + # `order_by("upper1")` means; + # * `self._fields` -- re-selecting restarts the counter, so a second + # `select(Upper(...))` would regenerate the first one's alias and + # collide with it. + taken = {f for f in fields if not isinstance(f, ResolvableExpression)} + taken |= set(self.sql_query.annotations) + # A model is free to have a column literally named `upper1` or `f1`, + # which an alias must not shadow even on a first select. + taken |= {f.name for f in self.model._model_meta.get_fields()} + if self._fields: + taken |= set(self._fields) + _fields = [] expressions = {} counter = 1 @@ -1865,8 +1939,9 @@ def values_list(self, *fields: str, flat: bool = False) -> QuerySet[Any]: while True: field_id = field_id_prefix + str(counter) counter += 1 - if field_id not in field_names: + if field_id not in taken: break + taken.add(field_id) expressions[field_id] = field _fields.append(field_id) else: @@ -1876,6 +1951,209 @@ def values_list(self, *fields: str, flat: bool = False) -> QuerySet[Any]: clone._iterable_class = FlatValuesListIterable if flat else ValuesListIterable return clone + # ---- select(): typed column selection returning honest rows ---- + # + # The ladder unwraps each Selectable[T] argument to its T and reassembles + # the row type. The precise row rides on RowQuerySet[R], a QuerySet flavor + # whose iteration yields R instead of model instances. A field binds its + # real value type; an expression is Selectable[Any], so it contributes Any + # while the fields around it stay precise. + + @overload + def select[S]( + self, item: Selectable[S], /, *, flat: Literal[True] + ) -> RowQuerySet[S]: ... + + @overload + def select[D]( + self, *items: Selectable[Any], result_type: type[D] + ) -> RowQuerySet[D]: ... + + @overload + def select[T0]( + self, i0: Selectable[T0], /, *, flat: Literal[False] = False + ) -> RowQuerySet[tuple[T0]]: ... + + @overload + def select[T0, T1]( + self, i0: Selectable[T0], i1: Selectable[T1], /, *, flat: Literal[False] = False + ) -> RowQuerySet[tuple[T0, T1]]: ... + + @overload + def select[T0, T1, T2]( + self, + i0: Selectable[T0], + i1: Selectable[T1], + i2: Selectable[T2], + /, + *, + flat: Literal[False] = False, + ) -> RowQuerySet[tuple[T0, T1, T2]]: ... + + @overload + def select[T0, T1, T2, T3]( + self, + i0: Selectable[T0], + i1: Selectable[T1], + i2: Selectable[T2], + i3: Selectable[T3], + /, + *, + flat: Literal[False] = False, + ) -> RowQuerySet[tuple[T0, T1, T2, T3]]: ... + + @overload + def select[T0, T1, T2, T3, T4]( + self, + i0: Selectable[T0], + i1: Selectable[T1], + i2: Selectable[T2], + i3: Selectable[T3], + i4: Selectable[T4], + /, + *, + flat: Literal[False] = False, + ) -> RowQuerySet[tuple[T0, T1, T2, T3, T4]]: ... + + @overload + def select[T0, T1, T2, T3, T4, T5]( + self, + i0: Selectable[T0], + i1: Selectable[T1], + i2: Selectable[T2], + i3: Selectable[T3], + i4: Selectable[T4], + i5: Selectable[T5], + /, + *, + flat: Literal[False] = False, + ) -> RowQuerySet[tuple[T0, T1, T2, T3, T4, T5]]: ... + + @overload + def select[T0, T1, T2, T3, T4, T5, T6]( + self, + i0: Selectable[T0], + i1: Selectable[T1], + i2: Selectable[T2], + i3: Selectable[T3], + i4: Selectable[T4], + i5: Selectable[T5], + i6: Selectable[T6], + /, + *, + flat: Literal[False] = False, + ) -> RowQuerySet[tuple[T0, T1, T2, T3, T4, T5, T6]]: ... + + @overload + def select[T0, T1, T2, T3, T4, T5, T6, T7]( + self, + i0: Selectable[T0], + i1: Selectable[T1], + i2: Selectable[T2], + i3: Selectable[T3], + i4: Selectable[T4], + i5: Selectable[T5], + i6: Selectable[T6], + i7: Selectable[T7], + /, + *, + flat: Literal[False] = False, + ) -> RowQuerySet[tuple[T0, T1, T2, T3, T4, T5, T6, T7]]: ... + + @overload + def select[T0, T1, T2, T3, T4, T5, T6, T7, T8]( + self, + i0: Selectable[T0], + i1: Selectable[T1], + i2: Selectable[T2], + i3: Selectable[T3], + i4: Selectable[T4], + i5: Selectable[T5], + i6: Selectable[T6], + i7: Selectable[T7], + i8: Selectable[T8], + /, + *, + flat: Literal[False] = False, + ) -> RowQuerySet[tuple[T0, T1, T2, T3, T4, T5, T6, T7, T8]]: ... + + @overload + def select[T0, T1, T2, T3, T4, T5, T6, T7, T8, T9]( + self, + i0: Selectable[T0], + i1: Selectable[T1], + i2: Selectable[T2], + i3: Selectable[T3], + i4: Selectable[T4], + i5: Selectable[T5], + i6: Selectable[T6], + i7: Selectable[T7], + i8: Selectable[T8], + i9: Selectable[T9], + /, + *, + flat: Literal[False] = False, + ) -> RowQuerySet[tuple[T0, T1, T2, T3, T4, T5, T6, T7, T8, T9]]: ... + + @overload + def select( + self, *items: Selectable[Any], flat: Literal[False] = False + ) -> RowQuerySet[tuple[Any, ...]]: ... + + def select( + self, + *items: Selectable[Any], + flat: bool = False, + result_type: type | None = None, + ) -> Any: + if not items: + raise TypeError("select() requires at least one column to select.") + if flat and len(items) > 1: + raise TypeError( + f"select(flat=True) takes exactly one column, got {len(items)}." + ) + if flat and result_type is not None: + raise TypeError("select() cannot combine flat=True with result_type=.") + if not isinstance(self, RowQuerySet) and self._fields is not None: + raise TypeError("Cannot call select() after values() or values_list().") + if self._returning_fields is not None or self._returning_instances: + raise TypeError( + "Cannot call select() after returning() — returning() captures " + "the rows a write touched, and a select() queryset cannot " + "write. Drop the returning() call." + ) + if self._prefetch_related_lookups: + # A prefetch hangs related objects off each result's attributes, + # and a row -- tuple, scalar or dataclass -- has nowhere to put + # them. Left alone it is silently wasted work for tuples and an + # AttributeError for result_type=. + raise TypeError( + "Cannot call select() after prefetch_related() — prefetched " + "objects are attached to model instances, and select() returns " + "rows. Select the columns you need from the related model " + "instead." + ) + + dataclass_type: type[DataclassInstance] | None = None + if result_type is not None: + if not ( + isinstance(result_type, type) and dataclasses.is_dataclass(result_type) + ): + raise TypeError("select(result_type=...) requires a dataclass.") + dataclass_type = result_type + _check_result_type_matches(dataclass_type, items) + + for item in items: + self._check_column_model(item) + columns = [_selectable_to_column(item) for item in items] + + clone = self._values_list(tuple(columns), flat=flat) + clone.__class__ = RowQuerySet + clone._select_result_type = dataclass_type + if dataclass_type is not None: + clone._iterable_class = SelectDataclassIterable + return clone + def none(self) -> Self: """Return an empty QuerySet.""" clone = self._chain() @@ -1925,6 +2203,36 @@ def where(self, *conditions: Q) -> Self: self._check_condition_model(condition) return self.filter(*conditions) + def _check_column_model(self, item: Selectable[Any]) -> None: + """Reject a column built from another model's fields. + + `where()`'s problem exactly, one method along: `Field[T]` carries no + model identity, so `Order.query.select(User.email)` type-checks and + the name `"email"` then resolves against `Order` -- silently the wrong + column when both models have one, a `FieldError` from the compiler + when they don't. + + A traversed column reports the model its traversal started from, so + this fires before the traversal refusal does: being another model's + column is the root mistake, and "select columns on the queried model" + would be advice that doesn't help. + + Expressions carry no origin and are left alone -- `F("email")` and + `Upper("email")` are strings resolved against whatever query they land + in, the same as `filter()`'s kwargs. + """ + if not isinstance(item, Field): + return + source_model = item.source_model + if source_model is not None and source_model is not self.model: + raise TypeError( + f"select() got a column built from " + f"{source_model.__name__}.{item.name}, but this is a " + f"{self.model.__name__} queryset. Select " + f"{self.model.__name__}'s own field, or traverse to it from " + f"{self.model.__name__}." + ) + def _check_condition_model(self, condition: Q) -> None: """Reject a condition built from another model's fields. @@ -2075,6 +2383,16 @@ def annotate(self, *args: Any, **kwargs: Any) -> Self: Return a query set in which the returned objects have been annotated with extra data or aggregations. """ + return self._annotate(*args, **kwargs) + + def _annotate(self, *args: Any, **kwargs: Any) -> Self: + """The mechanism behind `annotate()`. + + Separate from the public method because `_values_list` annotates + internally to alias expression columns, and `RowQuerySet` refuses the + public `annotate()` -- selecting an expression twice must not trip a + guard aimed at callers adding a column to a finished row. + """ self._validate_values_are_expressions( args + tuple(kwargs.values()), method_name="annotate" ) @@ -2093,14 +2411,18 @@ def annotate(self, *args: Any, **kwargs: Any) -> Self: annotations.update(kwargs) clone = self._chain() + # On a row-mode queryset the selected columns are what an alias can + # collide with; otherwise it is the model's own fields. names = self._fields + conflicts_with = "a selected column" if names is None: names = {field.name for field in self.model._model_meta.get_fields()} + conflicts_with = "a field on the model" for alias, annotation in annotations.items(): if alias in names: raise ValueError( - f"The annotation '{alias}' conflicts with a field on the model." + f"The annotation '{alias}' conflicts with {conflicts_with}." ) clone.sql_query.add_annotation(annotation, alias) if clone.sql_query.lock_mode and ( @@ -2293,6 +2615,7 @@ def _clone(self) -> Self: c._known_related_objects = self._known_related_objects c._iterable_class = self._iterable_class c._fields = self._fields + c._select_result_type = self._select_result_type c._returning_fields = self._returning_fields c._returning_instances = self._returning_instances return c @@ -2329,14 +2652,43 @@ def _next_is_sticky(self) -> Self: return self def _merge_sanity_check(self, other: QuerySet[T]) -> None: - """Check that two QuerySet classes may be merged.""" - if self._fields is not None and ( - set(self.sql_query.values_select) != set(other.sql_query.values_select) + """Check that two QuerySet classes may be merged. + + Either side being in row mode is enough to matter: merging a row-mode + queryset with a model-mode one produces a query neither side describes, + and left unchecked `model_qs | row_qs` recurses until the stack runs + out. The guard used to look only at `self`, so it caught the merge from + one side and not the other. + """ + if self._fields is None and other._fields is None: + return + + if ( + self._fields != other._fields + or set(self.sql_query.values_select) != set(other.sql_query.values_select) or set(self.sql_query.annotation_select) != set(other.sql_query.annotation_select) ): raise TypeError( - f"Merging '{self.__class__.__name__}' classes must involve the same values in each case." + f"Merging '{self.__class__.__name__}' and " + f"'{other.__class__.__name__}' classes must involve the same " + f"values in each case." + ) + + # Same columns is not the same rows: tuples, flat scalars and + # dataclasses all select identically and differ only in how each row + # is built. Merging two of them would quietly hand back whichever + # shape the left operand happened to carry. + if ( + self._iterable_class is not other._iterable_class + or self._select_result_type is not other._select_result_type + ): + raise TypeError( + f"Merging '{self.__class__.__name__}' and " + f"'{other.__class__.__name__}' classes must produce the same " + f"row shape: these select the same columns but build rows " + f"differently (tuple, flat scalar and result_type= rows are " + f"not interchangeable)." ) def _merge_known_related_objects(self, other: QuerySet[T]) -> None: @@ -2378,6 +2730,237 @@ def _validate_values_are_expressions( ) +# Why traversal is out: a column reached through a relation comes back over a +# join, so a nullable relation yields None where the traversed field's type says +# it can't. `Post.author.id` has the same problem as `Post.author.profile.city` +# — the FK column is NULL when the row has no author — so both are refused +# until select() can express that, and values_list() remains the way to spell it. +_NO_TRAVERSAL_IN_SELECT = ( + "a column reached through a relation is nullable in a way its type doesn't " + "say, so select() refuses it for now. Use values_list() with the lookup " + "path instead." +) + + +def _selectable_to_column(item: Selectable[Any]) -> str | ResolvableExpression: + """Turn a select() argument into something the values_list plumbing accepts. + + A field becomes its column name; an expression is passed through (the + plumbing auto-aliases it). Strings and FK traversal get their own error so + the message points at the real fix. + """ + # Local import: these pull in fields.related, which imports this module at + # load time (circular). + from plain.postgres.fields.related_descriptors import ForwardForeignKeyDescriptor + from plain.postgres.fields.related_typed import RelatedFieldRef + + if isinstance(item, str): + raise TypeError( + f"select() takes typed column references like User.email, not " + f"strings. Got {item!r}." + ) + if isinstance(item, RelatedFieldRef | ForwardForeignKeyDescriptor): + # The relation itself (`Post.author`) or an intermediate hop. Its key + # column is `Post.author.id`, which is refused for the same reason. + raise TypeError( + "select() takes columns, not relations, and a relation's key " + "column is not selectable yet either — " + _NO_TRAVERSAL_IN_SELECT + ) + if isinstance(item, Field): + if not item.name: + # A field read off a mixin class rather than a model: the mixin + # holds the declaration, and only the model it is mixed into has + # an attached, named copy. + raise TypeError( + f"select() got an unattached {type(item).__name__}. Reading a " + f"field off a mixin class gives the declaration, which has no " + f"name or column yet -- read it off the model that mixes it " + f"in instead." + ) + if item.is_lookup_reference: + # A traversed leaf: `Field.with_lookup_prefix` hands back the + # related model's field carrying the relation path as its name. + raise TypeError( + f"select() cannot select {item.name!r}: " + + _NO_TRAVERSAL_IN_SELECT + + f' Here that is values_list("{item.name}").' + ) + return item.name + # Matches what `_values_list` accepts, so anything values_list() can select + # — `F("x")` included — select() can select too. + if isinstance(item, ResolvableExpression): + return item + raise TypeError( + f"select() takes fields and expressions, got {type(item).__name__}." + ) + + +def _result_type_parameters( + result_type: type[DataclassInstance], +) -> tuple[inspect.Parameter, ...]: + """The constructor parameters a row is built from. + + Read off the signature rather than `dataclasses.fields()`, because the two + disagree in both directions: an `init=False` field is computed by the + dataclass and can't be passed, while an `InitVar` is a constructor + parameter that never appears in `fields()` at all. The signature is what + `result_type(*row)` actually has to satisfy. + """ + parameters = tuple(inspect.signature(result_type).parameters.values()) + for parameter in parameters: + if parameter.kind in ( + inspect.Parameter.VAR_POSITIONAL, + inspect.Parameter.VAR_KEYWORD, + ): + raise TypeError( + f"select(result_type={result_type.__name__}) needs a fixed " + f"constructor signature to map columns onto, but " + f"{result_type.__name__} takes {parameter!s}." + ) + return parameters + + +def _check_result_type_matches( + result_type: type[DataclassInstance], items: tuple[Selectable[Any], ...] +) -> None: + """Validate a dataclass result_type against the selected items. + + Arity must match the constructor, and each selected field's name must + equal the parameter at the same position — expressions are anonymous and + only need the position to line up. + """ + parameters = _result_type_parameters(result_type) + if len(parameters) != len(items): + raise TypeError( + f"select(result_type={result_type.__name__}) takes " + f"{len(parameters)} constructor arguments but {len(items)} " + f"columns were selected." + ) + for item, parameter in zip(items, parameters, strict=True): + if isinstance(item, Field) and item.name != parameter.name: + raise TypeError( + f"select(result_type={result_type.__name__}) maps columns " + f"positionally: field {item.name!r} does not match dataclass " + f"field {parameter.name!r} at the same position." + ) + + +class RowQuerySet[R](QuerySet[Any]): + """A QuerySet in row mode, returned by `select()`. + + Iteration yields the selected row type `R` — a tuple, a scalar (flat), or a + dataclass (result_type) — never a model instance. Everything that reads + rows is inherited: the parent's values_list machinery builds the SQL and + the rows, and `update()`/`delete()` already refuse a row-mode queryset. + + All this class adds at runtime is the refusal of the methods that would + re-enter row mode or hand back a row where a model instance is promised. + The rest is the `R` that the base, typed `QuerySet[Any]`, can't carry — + declared for the checker only, so row iteration doesn't pay for a Python + frame per call. + """ + + if TYPE_CHECKING: + + def __iter__(self) -> Iterator[R]: ... + + def first(self) -> R | None: ... + + def last(self) -> R | None: ... + + def get(self, *args: Any, **kwargs: Any) -> R: ... + + def get_or_none(self, *args: Any, **kwargs: Any) -> R | None: ... + + def iterator(self, chunk_size: int | None = None) -> Iterator[R]: ... + + @overload + def __getitem__(self, k: int) -> R: ... + + @overload + def __getitem__(self, k: slice) -> RowQuerySet[R]: ... + + # `Never` doesn't reject the call itself — the TypeError does that — but it + # does tell the checker control never returns, so a caller's trailing code + # reads as unreachable rather than as a QuerySet or a model instance. + + def annotate(self, *args: Any, **kwargs: Any) -> Never: + # An annotation appends a column, so the rows would gain a member the + # declared R doesn't have — silently for tuples, as a confusing + # constructor error for result_type=, and silently dropped for flat. + raise TypeError( + "Cannot call annotate() after select() — an annotation adds a " + "column, which would change the row shape out from under the " + "selected type. Annotate first, then select()." + ) + + def prefetch_related(self, *lookups: str | Prefetch | None) -> Never: + raise TypeError( + "Cannot call prefetch_related() after select() — prefetched " + "objects are attached to model instances, and select() returns " + "rows. Select the columns you need from the related model instead." + ) + + def create(self, **kwargs: Any) -> Never: + raise TypeError("Cannot call create() on a select() queryset.") + + def bulk_create(self, *args: Any, **kwargs: Any) -> Never: + raise TypeError("Cannot call bulk_create() on a select() queryset.") + + def values(self, *fields: str, **expressions: Any) -> Never: + raise TypeError("Cannot call values() after select().") + + def values_list(self, *fields: str, flat: bool = False) -> Never: + raise TypeError("Cannot call values_list() after select().") + + def get_or_create( + self, defaults: dict[str, Any] | None = None, **kwargs: Any + ) -> Never: + # The base would hand back whatever get() returns on a hit — a row — + # and a model instance on a miss. + raise TypeError("Cannot call get_or_create() after select().") + + def upsert( + self, + *_positional: Never, + defaults: dict[str, Any] | None = None, + create_defaults: dict[str, Any] | None = None, + conflict_defaults: dict[str, Any] | None = None, + unique_fields: Sequence[Field[Any] | type[Model]], + **kwargs: Any, + ) -> Never: + raise TypeError("Cannot call upsert() after select().") + + def bulk_upsert( + self, + objs: Sequence[Any], + *, + update_fields: Sequence[Field[Any] | type[Model]], + unique_fields: Sequence[Field[Any] | type[Model]], + batch_size: int | None = None, + ) -> Never: + raise TypeError("Cannot call bulk_upsert() after select().") + + def bulk_update( + self, objs: Sequence[Any], fields: list[str], batch_size: int | None = None + ) -> Never: + # The base would reach the same refusal, but only from the update() + # inside its own `transaction.atomic(savepoint=False)` -- which leaves + # the enclosing transaction unusable. Refusing up front keeps the + # failure a plain TypeError. + raise TypeError("Cannot call bulk_update() on a select() queryset.") + + def returning(self, *fields: Field[Any]) -> Never: + # returning() captures the rows a write touched, and select() has + # already refused every write. Accepting it would be inert at runtime + # and a lie statically -- the checker would believe update() hands + # back instances. + raise TypeError( + "Cannot call returning() after select() — returning() captures the " + "rows a write touched, and a select() queryset cannot write." + ) + + if TYPE_CHECKING: class ReturningQuerySet[T: "Model", R](QuerySet[T]): diff --git a/plain-postgres/plain/postgres/selectable.py b/plain-postgres/plain/postgres/selectable.py new file mode 100644 index 0000000000..eafb28941a --- /dev/null +++ b/plain-postgres/plain/postgres/selectable.py @@ -0,0 +1,20 @@ +"""The common base for anything `select()` can pull back as a column. + +`Field[T]` carries a real value type and subclasses `Selectable[T]`, so a +field contributes its `T` to the selected row. Expressions subclass +`Selectable[Any]` — they stay un-generic for now and contribute `Any`, which is +honest (an aggregate's output type isn't tracked yet) and can be tightened later +without touching `select()`. + +The overload ladder on `QuerySet.select()` binds each argument's `T` through +this shared base: a checker solves `T0` from `Field[str] <: Selectable[T0]` by +specializing the base class, so the marker needs no members of its own. The +ladder's per-column types are asserted in `tests/typing/select_rows.py`, which +is what would catch a checker that can't. +""" + +from __future__ import annotations + + +class Selectable[T]: + pass diff --git a/plain-postgres/plain/postgres/sql/datastructures.py b/plain-postgres/plain/postgres/sql/datastructures.py index 7856ec11e9..315d9349f4 100644 --- a/plain-postgres/plain/postgres/sql/datastructures.py +++ b/plain-postgres/plain/postgres/sql/datastructures.py @@ -32,10 +32,6 @@ def __init__( self.names_with_path = path_with_names -class Empty: - pass - - class Join: """ Used by sql.Query and sql.SQLCompiler to generate JOIN clauses into the diff --git a/plain-postgres/plain/postgres/sql/query.py b/plain-postgres/plain/postgres/sql/query.py index 0219682b42..5904813e7c 100644 --- a/plain-postgres/plain/postgres/sql/query.py +++ b/plain-postgres/plain/postgres/sql/query.py @@ -55,7 +55,7 @@ refs_expression, ) from plain.postgres.sql.constants import INNER, LOUTER, ORDER_DIR, SINGLE -from plain.postgres.sql.datastructures import BaseTable, Empty, Join, MultiJoin +from plain.postgres.sql.datastructures import BaseTable, Join, MultiJoin from plain.postgres.sql.where import AND, OR, NothingNode, WhereNode from plain.utils.regex_helper import _lazy_re_compile @@ -323,9 +323,7 @@ def clone(self) -> Self: Return a copy of the current Query. A lightweight alternative to deepcopy(). """ - obj = Empty() - obj.__class__ = self.__class__ - obj = cast(Self, obj) # Type checker doesn't understand __class__ reassignment + obj = object.__new__(self.__class__) # Copy references to everything. obj.__dict__ = self.__dict__.copy() # Clone attributes that can't use shallow copy. diff --git a/plain-postgres/tests/app/examples/migrations/0028_aliascollisionexample.py b/plain-postgres/tests/app/examples/migrations/0028_aliascollisionexample.py new file mode 100644 index 0000000000..9eb3aecd5b --- /dev/null +++ b/plain-postgres/tests/app/examples/migrations/0028_aliascollisionexample.py @@ -0,0 +1,24 @@ +# Generated by Plain 0.163.1 on 2026-09-20 02:36 + +from plain.postgres import migrations + +from plain import postgres + + +class Migration(migrations.Migration): + dependencies = (("examples", "0027_upsertstamped"),) + + operations = ( + migrations.CreateModel( + name="AliasCollisionExample", + fields=[ + ("id", postgres.PrimaryKeyField()), + ("f1", postgres.TextField(default="", max_length=100, required=False)), + ("name", postgres.TextField(max_length=100)), + ( + "upper1", + postgres.TextField(default="", max_length=100, required=False), + ), + ], + ), + ) diff --git a/plain-postgres/tests/app/examples/models/__init__.py b/plain-postgres/tests/app/examples/models/__init__.py index a79bafaba8..ace103d459 100644 --- a/plain-postgres/tests/app/examples/models/__init__.py +++ b/plain-postgres/tests/app/examples/models/__init__.py @@ -2,6 +2,7 @@ # Import submodules so @postgres.register_model runs for every test model. from . import ( # noqa: F401 + alias_collisions, constraints, defaults, delete, diff --git a/plain-postgres/tests/app/examples/models/alias_collisions.py b/plain-postgres/tests/app/examples/models/alias_collisions.py new file mode 100644 index 0000000000..94cbea7246 --- /dev/null +++ b/plain-postgres/tests/app/examples/models/alias_collisions.py @@ -0,0 +1,19 @@ +from __future__ import annotations + +from plain.postgres import Field, types + +from plain import postgres + + +@postgres.register_model +class AliasCollisionExample(postgres.Model): + """A model whose column names look like generated expression aliases. + + `select(Upper(...))` names its column `upper1` and `select(F(...))` names + its column `f1`, so a model that already has columns by those names is + what proves the generator skips them rather than shadowing a real column. + """ + + name: Field[str] = types.TextField(max_length=100) + upper1: Field[str] = types.TextField(max_length=100, required=False, default="") + f1: Field[str] = types.TextField(max_length=100, required=False, default="") diff --git a/plain-postgres/tests/public/test_select.py b/plain-postgres/tests/public/test_select.py new file mode 100644 index 0000000000..1c1b1577d7 --- /dev/null +++ b/plain-postgres/tests/public/test_select.py @@ -0,0 +1,914 @@ +"""Runtime behavior of `QuerySet.select()` — typed column selection that +returns honest rows (tuples, scalars, or dataclasses), never partial model +instances. + +The static-typing contract lives in tests/typing/select_rows.py. +""" + +from __future__ import annotations + +from dataclasses import InitVar, dataclass, field, fields + +import pytest +from app.examples.models.alias_collisions import AliasCollisionExample +from app.examples.models.defaults import DefaultsExample +from app.examples.models.mixins import MixinTestModel, TimestampMixin +from app.examples.models.relationships import Tag, Widget, WidgetTag +from plain.postgres import RowQuerySet +from plain.postgres.aggregates import Count +from plain.postgres.expressions import F, Value +from plain.postgres.functions import Lower, Upper + + +@pytest.fixture +def rows(db): + DefaultsExample.query.create(name="alpha", priority=3, note="first") + DefaultsExample.query.create(name="beta", priority=1, note=None) + DefaultsExample.query.create(name="gamma", priority=2, note="third") + + +def test_select_returns_tuple_rows(rows): + result = ( + DefaultsExample.query.order_by("priority") + .select(DefaultsExample.name, DefaultsExample.priority) + .all() + ) + assert list(result) == [("beta", 1), ("gamma", 2), ("alpha", 3)] + + +def test_select_flat_returns_scalars(rows): + names = ( + DefaultsExample.query.order_by("name") + .select(DefaultsExample.name, flat=True) + .all() + ) + assert list(names) == ["alpha", "beta", "gamma"] + + +def test_select_single_column_is_one_tuple(rows): + result = DefaultsExample.query.order_by("name").select(DefaultsExample.name) + assert list(result) == [("alpha",), ("beta",), ("gamma",)] + + +def test_select_preserves_nullable_values(rows): + result = ( + DefaultsExample.query.order_by("name") + .select(DefaultsExample.name, DefaultsExample.note) + .all() + ) + assert list(result) == [("alpha", "first"), ("beta", None), ("gamma", "third")] + + +def test_select_chains_with_where(rows): + result = DefaultsExample.query.where(DefaultsExample.priority.gte(2)).select( + DefaultsExample.name, flat=True + ) + assert sorted(result) == ["alpha", "gamma"] + + +def test_select_expression_column(rows): + result = ( + DefaultsExample.query.order_by("name") + .select(DefaultsExample.priority, Upper("name")) + .all() + ) + assert list(result) == [(3, "ALPHA"), (1, "BETA"), (2, "GAMMA")] + + +def test_select_f_expression_column(rows): + """F() is a Combinable, not a BaseExpression -- select() takes it anyway, + because values_list() does. Static half: tests/typing/select_rows.py.""" + result = ( + DefaultsExample.query.order_by("name") + .select(DefaultsExample.name, F("priority")) + .all() + ) + assert list(result) == [("alpha", 3), ("beta", 1), ("gamma", 2)] + + +def test_select_combined_f_expression_column(rows): + result = DefaultsExample.query.order_by("name").select(F("priority") + 1, flat=True) + assert list(result) == [4, 2, 3] + + +def test_select_returns_row_queryset(rows): + result = DefaultsExample.query.select(DefaultsExample.name) + assert isinstance(result, RowQuerySet) + + +def test_select_first_returns_row(rows): + row = ( + DefaultsExample.query.order_by("priority") + .select(DefaultsExample.name, DefaultsExample.priority) + .first() + ) + assert row == ("beta", 1) + + +def test_select_get_returns_row(rows): + row = ( + DefaultsExample.query.where(DefaultsExample.name.equals("alpha")) + .select(DefaultsExample.name, DefaultsExample.priority) + .get() + ) + assert row == ("alpha", 3) + + +# ---- result_type=Dataclass ---- + + +@dataclass +class NameStat: + name: str + priority: int + + +def test_select_result_type_builds_dataclass(rows): + result = ( + DefaultsExample.query.order_by("priority") + .select(DefaultsExample.name, DefaultsExample.priority, result_type=NameStat) + .all() + ) + assert list(result) == [ + NameStat(name="beta", priority=1), + NameStat(name="gamma", priority=2), + NameStat(name="alpha", priority=3), + ] + + +def test_select_result_type_with_expression_column(rows): + @dataclass + class NameUpper: + priority: int + upper: str + + result = ( + DefaultsExample.query.order_by("name") + .select( + DefaultsExample.priority, + Upper("name"), + result_type=NameUpper, + ) + .all() + ) + assert result[0] == NameUpper(priority=3, upper="ALPHA") + + +# ---- error paths ---- + + +def test_select_requires_at_least_one_column(db): + with pytest.raises(TypeError, match="at least one column"): + DefaultsExample.query.select() + + +def test_select_rejects_string_argument(db): + with pytest.raises(TypeError, match="strings"): + DefaultsExample.query.select("name") # ty: ignore[no-matching-overload] + + +def test_select_rejects_fk_traversal(db): + """The message names the values_list() spelling that does work.""" + with pytest.raises(TypeError, match=r'values_list\("widget__name"\)'): + WidgetTag.query.select(WidgetTag.widget.name) + + +def test_select_rejects_fk_reference(db): + """A relation is not a column, and its key column is out for now too.""" + with pytest.raises(TypeError, match="key column is not selectable yet"): + WidgetTag.query.select(WidgetTag.widget) # ty: ignore[no-matching-overload] + + +def test_select_flat_rejects_more_than_one_column(db): + """The message names select(), not the values_list() plumbing underneath.""" + with pytest.raises(TypeError, match=r"select\(flat=True\) takes exactly one"): + DefaultsExample.query.select( # ty: ignore[no-matching-overload] + DefaultsExample.name, DefaultsExample.priority, flat=True + ) + + +def test_select_rejects_the_foreign_key_column_for_now(db): + """`WidgetTag.widget.id` is the FK column, but it still arrives over the + relation, so it is refused with the rest of traversal until select() can + say it is nullable.""" + with pytest.raises(TypeError, match=r'values_list\("widget__id"\)'): + WidgetTag.query.select(WidgetTag.widget.id) + + +def test_select_result_type_must_be_dataclass(db): + with pytest.raises(TypeError, match="dataclass"): + DefaultsExample.query.select(DefaultsExample.name, result_type=dict) + + +def test_select_result_type_arity_must_match(db): + with pytest.raises(TypeError, match="constructor arguments"): + DefaultsExample.query.select(DefaultsExample.name, result_type=NameStat) + + +def test_select_result_type_field_name_must_match(db): + @dataclass + class Renamed: + label: str + priority: int + + with pytest.raises(TypeError, match="positionally"): + DefaultsExample.query.select( + DefaultsExample.name, DefaultsExample.priority, result_type=Renamed + ) + + +def test_select_flat_and_result_type_conflict(db): + with pytest.raises(TypeError, match="combine flat"): + DefaultsExample.query.select( # ty: ignore[no-matching-overload] + DefaultsExample.name, flat=True, result_type=NameStat + ) + + +def test_select_result_type_ignores_init_false_fields(rows): + """An init=False field is computed by the dataclass, so it is neither + selected nor counted against the arity.""" + + @dataclass + class Computed: + name: str + priority: int + label: str = field(default="", init=False) + + def __post_init__(self) -> None: + self.label = f"{self.name}:{self.priority}" + + result = ( + DefaultsExample.query.order_by("name") + .select(DefaultsExample.name, DefaultsExample.priority, result_type=Computed) + .first() + ) + assert result == Computed(name="alpha", priority=3) + assert result.label == "alpha:3" + + +def test_select_result_type_counts_init_vars(rows): + """An InitVar is a constructor parameter that never appears in + dataclasses.fields(), so the columns are mapped off the signature.""" + + @dataclass + class WithInitVar: + name: str + priority: InitVar[int] + + def __post_init__(self, priority: int) -> None: + # `priority` is a constructor argument only -- consumed here and + # never stored, which is why fields() doesn't list it. + assert priority > 0 + + assert [f.name for f in fields(WithInitVar)] == ["name"] + + result = ( + DefaultsExample.query.order_by("name") + .select(DefaultsExample.name, DefaultsExample.priority, result_type=WithInitVar) + .first() + ) + assert result == WithInitVar(name="alpha", priority=3) + + +def test_select_result_type_handles_init_var_and_init_false_together(rows): + """fields() disagrees with the constructor in both directions here: it + lists `computed` (which can't be passed) and omits `priority` (which + must be). The signature is right on both counts.""" + + @dataclass + class Mixed: + name: str + priority: InitVar[int] + computed: str = field(default="", init=False) + + def __post_init__(self, priority: int) -> None: + self.computed = f"{self.name}/{priority}" + + result = ( + DefaultsExample.query.order_by("name") + .select(DefaultsExample.name, DefaultsExample.priority, result_type=Mixed) + .first() + ) + assert result == Mixed(name="alpha", priority=3) + assert result.computed == "alpha/3" + + +def test_select_result_type_arity_counts_init_vars(rows): + @dataclass + class OnlyInitVar: + name: str + priority: InitVar[int] + + def __post_init__(self, priority: int) -> None: + pass + + with pytest.raises(TypeError, match="takes 2 constructor arguments but 1"): + DefaultsExample.query.select(DefaultsExample.name, result_type=OnlyInitVar) + + +class TestResultTypeParameterKinds: + """Parameter kind decides how each value has to be passed. A signature + always orders positionals before keyword-onlys, so the row splits at one + point -- but the split has to happen, because a positional-only parameter + can't be filled by name and a keyword-only one can't be filled by + position.""" + + def test_all_positional(self, rows): + @dataclass + class AllPositional: + name: str + priority: int + + result = ( + DefaultsExample.query.order_by("name") + .select( + DefaultsExample.name, + DefaultsExample.priority, + result_type=AllPositional, + ) + .first() + ) + assert result == AllPositional(name="alpha", priority=3) + + def test_all_keyword_only(self, rows): + @dataclass(kw_only=True) + class AllKeyword: + name: str + priority: int + + result = ( + DefaultsExample.query.order_by("name") + .select( + DefaultsExample.name, DefaultsExample.priority, result_type=AllKeyword + ) + .first() + ) + assert result == AllKeyword(name="alpha", priority=3) + + def test_positional_only_mixed_with_keyword_only(self, rows): + """Neither the all-positional nor the all-keyword call works here.""" + + @dataclass + class Mixed: + name: str + priority: int + + def __init__(self, name: str, /, *, priority: int) -> None: + self.name = name + self.priority = priority + + result = ( + DefaultsExample.query.order_by("name") + .select(DefaultsExample.name, DefaultsExample.priority, result_type=Mixed) + .first() + ) + assert result == Mixed("alpha", priority=3) + + def test_positional_or_keyword_mixed_with_keyword_only(self, rows): + @dataclass + class PartlyKwOnly: + name: str + priority: int = field(kw_only=True) + + result = ( + DefaultsExample.query.order_by("name") + .select( + DefaultsExample.name, DefaultsExample.priority, result_type=PartlyKwOnly + ) + .first() + ) + assert result == PartlyKwOnly(name="alpha", priority=3) + + +def test_select_result_type_rejects_a_variadic_constructor(rows): + @dataclass(init=False) + class Variadic: + name: str + + def __init__(self, *args: object) -> None: + self.name = str(args[0]) + + with pytest.raises(TypeError, match="fixed constructor signature"): + DefaultsExample.query.select(DefaultsExample.name, result_type=Variadic) + + +class TestSelectTwiceIsAlwaysLastWins: + """Re-selecting replaces the columns, whatever either side is made of. + + The internal alias a selected expression gets used to restart its counter + from 1 on every select(), so a second expression column regenerated the + first one's alias and collided with it -- a ValueError blaming "a field on + the model" that named neither the real cause nor a real field. + """ + + @pytest.fixture + def ordered(self, db): + # LOWER(name) and UPPER(status) disagree about order, so a test can + # tell which one a query actually used. + DefaultsExample.query.create(name="zzz", priority=1, status="aaa") + DefaultsExample.query.create(name="aaa", priority=2, status="zzz") + + def test_field_then_field(self, ordered): + result = ( + DefaultsExample.query.order_by("name") + .select(DefaultsExample.name) + .select(DefaultsExample.priority) + ) + assert list(result) == [(2,), (1,)] + + def test_field_then_expression(self, ordered): + result = ( + DefaultsExample.query.order_by("name") + .select(DefaultsExample.name) + .select(Upper("name")) + ) + assert list(result) == [("AAA",), ("ZZZ",)] + + def test_expression_then_field(self, ordered): + result = ( + DefaultsExample.query.order_by("name") + .select(Upper("name")) + .select(DefaultsExample.name) + ) + assert list(result) == [("aaa",), ("zzz",)] + + def test_expression_then_same_expression_function(self, ordered): + """The reported crash: both aliases wanted to be 'upper1'.""" + result = ( + DefaultsExample.query.order_by("name") + .select(Upper("name")) + .select(Upper("status")) + ) + assert list(result) == [("ZZZ",), ("AAA",)] + + def test_expression_then_different_expression_function(self, ordered): + result = ( + DefaultsExample.query.order_by("name") + .select(Upper("name")) + .select(Lower("status")) + ) + assert list(result) == [("zzz",), ("aaa",)] + + def test_f_then_f(self, ordered): + result = ( + DefaultsExample.query.order_by("name") + .select(F("priority")) + .select(F("name")) + ) + assert list(result) == [("aaa",), ("zzz",)] + + def test_mixed_list_then_mixed_list(self, ordered): + result = ( + DefaultsExample.query.order_by("name") + .select(DefaultsExample.name, Upper("status")) + .select(DefaultsExample.priority, Lower("name")) + ) + assert list(result) == [(2, "aaa"), (1, "zzz")] + + def test_flat_then_flat(self, ordered): + result = ( + DefaultsExample.query.order_by("name") + .select(Upper("name"), flat=True) + .select(Upper("status"), flat=True) + ) + assert list(result) == ["ZZZ", "AAA"] + + def test_expression_then_result_type(self, ordered): + @dataclass + class NameAndUpper: + name: str + upper: str + + result = ( + DefaultsExample.query.order_by("name") + .select(Upper("name")) + .select(DefaultsExample.name, Upper("status"), result_type=NameAndUpper) + ) + assert list(result) == [ + NameAndUpper(name="aaa", upper="ZZZ"), + NameAndUpper(name="zzz", upper="AAA"), + ] + + def test_three_selects_in_a_row(self, ordered): + result = ( + DefaultsExample.query.order_by("name") + .select(Upper("name")) + .select(Upper("status")) + .select(Upper("name")) + ) + assert list(result) == [("AAA",), ("ZZZ",)] + + def test_aggregates_still_work_after_re_selecting(self, ordered): + rows = DefaultsExample.query.select(Upper("name")).select(Upper("status")) + assert rows.count() == 2 + assert rows.exists() is True + + +class TestInternalAliasesNeverClobberUserAnnotations: + """A selected expression gets an internal alias, and that alias must not + land on one the caller already chose -- doing so silently redefines what + order_by()/filter() on that name mean.""" + + @pytest.fixture + def ordered(self, db): + DefaultsExample.query.create(name="zzz", priority=1, status="aaa") + DefaultsExample.query.create(name="aaa", priority=2, status="zzz") + + def test_order_by_a_user_annotation_keeps_its_meaning(self, ordered): + """'upper1' is what the generator would have produced for Upper(...), + so an unguarded generator overwrote LOWER(name) with UPPER(status) and + silently reversed the order.""" + result = ( + DefaultsExample.query.annotate(upper1=Lower("name")) + .order_by("upper1") + .select(DefaultsExample.name, Upper("status")) + ) + assert [row[0] for row in result] == ["aaa", "zzz"] + + def test_filter_on_a_user_annotation_keeps_its_meaning(self, ordered): + result = ( + DefaultsExample.query.annotate(upper1=Lower("name")) + .filter(upper1="aaa") + .select(DefaultsExample.name, Upper("status")) + ) + assert list(result) == [("aaa", "ZZZ")] + + +class TestColumnsBelongToTheirModel: + """`select()`'s half of the guard `where()` carries. + + `Field[T]` carries no model identity, so nothing stops one model's field + being handed to another model's `select()`. The name then resolves against + the queried model -- silently the wrong column when both have one. + `Widget` is the other model here because it also has a `name`. + """ + + def test_same_model_column_passes(self, rows): + result = DefaultsExample.query.order_by("name").select( + DefaultsExample.name, flat=True + ) + assert list(result) == ["alpha", "beta", "gamma"] + + def test_other_model_column_raises_although_the_name_exists(self, rows): + """The dangerous case: both models have a `name`, so without the check + this quietly selected DefaultsExample.name.""" + with pytest.raises(TypeError) as excinfo: + DefaultsExample.query.select(Widget.name) + message = str(excinfo.value) + assert "Widget.name" in message + assert "DefaultsExample queryset" in message + + def test_other_model_column_raises_when_the_name_does_not_exist(self, rows): + """Previously a FieldError from deep in the compiler.""" + with pytest.raises(TypeError, match="Widget.size"): + DefaultsExample.query.select(Widget.size) + + def test_traversed_column_rooted_elsewhere_raises_as_cross_model(self, rows): + """Being another model's column is the root mistake, so it is named + ahead of the traversal refusal -- "select columns on the queried + model" would be advice that doesn't help here.""" + with pytest.raises(TypeError, match="WidgetTag.widget__name"): + DefaultsExample.query.select(WidgetTag.widget.name) + + def test_traversed_column_on_its_own_root_still_reports_traversal(self, db): + with pytest.raises(TypeError, match="reached through a relation"): + WidgetTag.query.select(WidgetTag.widget.name) + + def test_expressions_are_unaffected(self, rows): + """An expression takes a string resolved against whatever query it + lands in, like filter()'s kwargs -- there is no origin to check.""" + assert list( + DefaultsExample.query.order_by("name").select(F("priority"), flat=True) + ) == [3, 1, 2] + assert list( + DefaultsExample.query.order_by("name").select(Upper("name"), flat=True) + ) == ["ALPHA", "BETA", "GAMMA"] + + def test_mixed_list_with_one_foreign_column_raises(self, rows): + with pytest.raises(TypeError, match="Widget.name"): + DefaultsExample.query.select(DefaultsExample.priority, Widget.name) + + +def test_select_rejects_a_field_read_off_a_mixin(db): + """The mixin holds the declaration; only the model that mixes it in has an + attached, named copy. This used to surface as a bare AssertionError. + + A checker rejects the access too (`__get__` wants `owner: type[Model]`, + and a mixin isn't one), so this is the backstop for an untyped call site + rather than the only guard. + """ + with pytest.raises(TypeError, match="unattached"): + MixinTestModel.query.select(TimestampMixin.created_at) # ty: ignore[invalid-attribute-access] + + +def test_select_accepts_the_same_field_off_the_model(db): + MixinTestModel.query.create(name="a") + assert len(list(MixinTestModel.query.select(MixinTestModel.created_at))) == 1 + + +def test_select_alias_skips_a_column_that_looks_like_one(db): + """A model is free to have columns named `upper1` and `f1`; a generated + alias must step over them rather than shadow a real column.""" + AliasCollisionExample.query.create(name="a", upper1="real-upper1", f1="real-f1") + + result = AliasCollisionExample.query.select( + AliasCollisionExample.upper1, Upper("name") + ) + assert list(result) == [("real-upper1", "A")] + + result = AliasCollisionExample.query.select(AliasCollisionExample.f1, F("name")) + assert list(result) == [("real-f1", "a")] + + +class TestMergingRowAndModelQuerysets: + """Merging a row-mode queryset with a model-mode one produces a query + neither side describes. The guard only looked at the left operand, so + `model_qs | row_qs` recursed until the stack ran out.""" + + def test_model_or_row_raises(self, rows): + with pytest.raises(TypeError, match="must involve the same values"): + DefaultsExample.query.all() | DefaultsExample.query.select( + DefaultsExample.name + ) + + def test_row_or_model_raises(self, rows): + with pytest.raises(TypeError, match="must involve the same values"): + ( + DefaultsExample.query.select(DefaultsExample.name) + | DefaultsExample.query.all() + ) + + def test_model_and_row_raises(self, rows): + with pytest.raises(TypeError, match="must involve the same values"): + DefaultsExample.query.all() & DefaultsExample.query.select( + DefaultsExample.name + ) + + def test_sliced_row_queryset_merges(self, rows): + """A sliced left operand is re-expressed as an id subquery, which used + to call the public values() and hit select()'s own refusal.""" + r = DefaultsExample.query.order_by("name").select(DefaultsExample.name) + assert len(list(r[0:1] | r)) == 3 + assert len(list(r | r[0:1])) == 3 + + def test_different_row_shapes_do_not_merge(self, rows): + """Same columns, different rows: tuple, flat and result_type= select + identically and differ only in how each row is built, so a merge used + to hand back whichever shape the left operand carried.""" + + @dataclass + class NameOnly: + name: str + + tuples = DefaultsExample.query.select(DefaultsExample.name) + flat = DefaultsExample.query.select(DefaultsExample.name, flat=True) + dataclasses_ = DefaultsExample.query.select( + DefaultsExample.name, result_type=NameOnly + ) + + for left, right in ( + (dataclasses_, tuples), + (tuples, dataclasses_), + (tuples, flat), + (flat, tuples), + ): + with pytest.raises(TypeError, match="same row shape"): + left | right + + def test_matching_row_shapes_still_merge(self, rows): + @dataclass + class NameOnly: + name: str + + left = DefaultsExample.query.select(DefaultsExample.name, result_type=NameOnly) + right = DefaultsExample.query.select(DefaultsExample.name, result_type=NameOnly) + assert len(list(left | right)) == 3 + + def test_matching_sides_still_merge(self, rows): + model = DefaultsExample.query.all() | DefaultsExample.query.all() + assert len(list(model)) == 3 + row = DefaultsExample.query.select( + DefaultsExample.name + ) | DefaultsExample.query.select(DefaultsExample.name) + assert len(list(row)) == 3 + + +class TestSelectTwiceWithExpressions: + """`select()` twice is last-wins, and that has to hold for expression + columns too. `_values_list` aliases an expression by annotating + internally, which used to call the *public* annotate() and so trip + RowQuerySet's guard -- a guard meant for callers adding a column to a + finished row, not for select() rebuilding one.""" + + def test_replacing_fields_with_an_expression(self, rows): + result = ( + DefaultsExample.query.order_by("name") + .select(DefaultsExample.name) + .select(F("priority")) + ) + assert list(result) == [(3,), (1,), (2,)] + + def test_replacing_fields_with_a_flat_expression(self, rows): + result = ( + DefaultsExample.query.order_by("name") + .select(DefaultsExample.name) + .select(F("priority"), flat=True) + ) + assert list(result) == [3, 1, 2] + + def test_replacing_fields_with_an_expression_and_result_type(self, rows): + @dataclass + class NameAndExpr: + name: str + upper: str + + result = ( + DefaultsExample.query.order_by("name") + .select(DefaultsExample.priority) + .select(DefaultsExample.name, Upper("name"), result_type=NameAndExpr) + .first() + ) + assert result == NameAndExpr(name="alpha", upper="ALPHA") + + def test_annotate_is_still_refused_for_callers(self, rows): + """The guard the internal path now bypasses is still there.""" + with pytest.raises(TypeError, match="Annotate first, then select"): + DefaultsExample.query.select(DefaultsExample.name).annotate(x=Value(1)) + + +class TestPrefetchRelatedAndSelect: + """A prefetch hangs related objects off each result's attributes, and a + row has nowhere to put them: it was silently wasted work for tuples and + scalars, and an AttributeError for result_type=. Refused in both orders, + the same as select_related().""" + + @pytest.fixture + def widget(self, db): + w = Widget.query.create(name="w", size="s") + tag = Tag.query.create(name="t") + w.tags.add(tag) + return w + + def test_select_after_prefetch_raises_in_tuple_mode(self, widget): + with pytest.raises(TypeError, match="after prefetch_related"): + Widget.query.prefetch_related("tags").select(Widget.name) + + def test_select_after_prefetch_raises_in_flat_mode(self, widget): + with pytest.raises(TypeError, match="after prefetch_related"): + Widget.query.prefetch_related("tags").select(Widget.name, flat=True) + + def test_select_after_prefetch_raises_in_result_type_mode(self, widget): + @dataclass + class NameRow: + name: str + + with pytest.raises(TypeError, match="after prefetch_related"): + Widget.query.prefetch_related("tags").select( + Widget.name, result_type=NameRow + ) + + def test_prefetch_after_select_raises(self, widget): + with pytest.raises(TypeError, match="after select"): + Widget.query.select(Widget.name).prefetch_related("tags") + + def test_prefetch_without_select_is_unaffected(self, widget): + widgets = list(Widget.query.prefetch_related("tags")) + assert [t.name for t in widgets[0].tags.query.all()] == ["t"] + + +class TestAnnotateAfterSelect: + """An annotation appends a column, so it would change the row shape out + from under the type select() already declared. The supported order is + annotate first, then select().""" + + def test_annotate_after_select_raises(self, rows): + with pytest.raises(TypeError, match="Annotate first, then select"): + DefaultsExample.query.select(DefaultsExample.name).annotate(x=Value(1)) + + def test_annotate_after_select_raises_in_result_type_mode(self, rows): + with pytest.raises(TypeError, match="Annotate first, then select"): + DefaultsExample.query.select( + DefaultsExample.name, DefaultsExample.priority, result_type=NameStat + ).annotate(x=Value(1)) + + def test_annotate_after_select_raises_in_flat_mode(self, rows): + with pytest.raises(TypeError, match="Annotate first, then select"): + DefaultsExample.query.select(DefaultsExample.name, flat=True).annotate( + x=Value(1) + ) + + def test_annotate_before_select_still_works(self, rows): + """The supported order — the annotation is selectable as a column.""" + result = ( + DefaultsExample.query.annotate(n=Count("id")) + .order_by("name") + .select(DefaultsExample.name) + ) + assert list(result) == [("alpha",), ("beta",), ("gamma",)] + + def test_annotate_is_unaffected_on_a_plain_queryset(self, rows): + assert DefaultsExample.query.annotate(n=Count("id")).count() == 3 + + +def test_get_or_create_after_select_raises(db): + with pytest.raises(TypeError, match="get_or_create"): + DefaultsExample.query.select(DefaultsExample.name).get_or_create(name="x") + + +def test_bulk_update_after_select_raises_before_any_sql(db): + """The base would reach the same refusal, but only from the update() + inside its own `transaction.atomic(savepoint=False)` -- which leaves the + enclosing transaction unusable, so the *next* query fails too.""" + DefaultsExample.query.create(name="alpha", priority=3) + objs = list(DefaultsExample.query.all()) + for obj in objs: + obj.name = "changed" + + with pytest.raises(TypeError, match="bulk_update"): + DefaultsExample.query.select(DefaultsExample.name).bulk_update(objs, ["name"]) + + # The transaction is still usable: nothing was sent. + assert DefaultsExample.query.count() == 1 + assert DefaultsExample.query.get().name == "alpha" + + +def test_returning_after_select_raises(db): + with pytest.raises(TypeError, match="returning"): + DefaultsExample.query.select(DefaultsExample.name).returning() + + +def test_select_after_returning_raises(db): + with pytest.raises(TypeError, match="after returning"): + DefaultsExample.query.returning().select(DefaultsExample.name) + + +def test_returning_without_select_is_unaffected(db): + DefaultsExample.query.create(name="alpha", priority=3) + updated = DefaultsExample.query.returning().update(name="beta") + assert [obj.name for obj in updated] == ["beta"] + + +def test_upsert_after_select_raises(db): + with pytest.raises(TypeError, match="upsert"): + DefaultsExample.query.select(DefaultsExample.name).upsert( + name="x", unique_fields=[DefaultsExample.name] + ) + + +def test_bulk_upsert_after_select_raises(db): + with pytest.raises(TypeError, match="bulk_upsert"): + DefaultsExample.query.select(DefaultsExample.name).bulk_upsert( + [DefaultsExample(name="x")], + update_fields=[DefaultsExample.priority], + unique_fields=[DefaultsExample.name], + ) + + +def test_select_after_values_raises(db): + with pytest.raises(TypeError, match="after values"): + DefaultsExample.query.values("name").select(DefaultsExample.name) + + +def test_values_after_select_raises(db): + with pytest.raises(TypeError, match="after select"): + DefaultsExample.query.select(DefaultsExample.name).values("name") + + +def test_create_after_select_raises(db): + with pytest.raises(TypeError, match="create"): + DefaultsExample.query.select(DefaultsExample.name).create(name="x") + + +def test_bulk_create_after_select_raises(db): + with pytest.raises(TypeError, match="bulk_create"): + DefaultsExample.query.select(DefaultsExample.name).bulk_create( + [DefaultsExample(name="x")] + ) + + +def test_update_after_select_raises(db): + with pytest.raises(TypeError, match="update"): + DefaultsExample.query.select(DefaultsExample.name).update(name="x") + + +def test_delete_after_select_raises(db): + with pytest.raises(TypeError, match="delete"): + DefaultsExample.query.select(DefaultsExample.name).delete() + + +def test_update_refuses_any_row_mode_queryset(db): + """select() shares delete()'s guard rather than adding its own, so + update() now refuses values()/values_list() for the same reason.""" + with pytest.raises(TypeError, match="Cannot call update"): + DefaultsExample.query.values("name").update(name="x") + with pytest.raises(TypeError, match="Cannot call update"): + DefaultsExample.query.values_list("name").update(name="x") + + +def test_select_twice_last_wins(rows): + result = ( + DefaultsExample.query.order_by("name") + .select(DefaultsExample.name) + .select(DefaultsExample.priority, flat=True) + ) + assert list(result) == [3, 1, 2] diff --git a/plain-postgres/tests/typing/queryset_access.py b/plain-postgres/tests/typing/queryset_access.py index b72b7136ef..ffe4fd6ef1 100644 --- a/plain-postgres/tests/typing/queryset_access.py +++ b/plain-postgres/tests/typing/queryset_access.py @@ -19,3 +19,28 @@ def must_accept_class_access() -> None: def must_reject_instance_access(row: DefaultsExample) -> None: # Runtime half: tests/public/test_manager_assignment.py. _ = row.query # ty: ignore[invalid-attribute-access] + + +def must_accept_chaining_that_keeps_the_custom_queryset() -> None: + """A chaining method returns `Self`, not `QuerySet[T]`. + + Anything that clones through `self._chain()` hands back the same class at + runtime, so annotating it `QuerySet[T]` would erase a custom subclass -- + `CustomQuerySetModel.query.filter(...).get_custom()` would stop + type-checking even though it works. + """ + rows = CustomQuerySetModel.query.filter(name="a") + assert_type(rows, CustomQuerySet) + assert_type(rows.order_by("name"), CustomQuerySet) + assert_type(rows.reverse(), CustomQuerySet) + assert_type(rows.none(), CustomQuerySet) + assert_type(rows.distinct(), CustomQuerySet) + assert_type(rows.only("name"), CustomQuerySet) + assert_type(rows.defer("name"), CustomQuerySet) + assert_type(rows.for_update(), CustomQuerySet) + assert_type(rows[0:2], CustomQuerySet) + assert_type(rows & rows, CustomQuerySet) + # The custom method stays reachable through the whole chain. (It has no + # return annotation of its own, so only reachability is claimed here -- + # an unmarked line that started erroring would fail the build.) + rows.order_by("name").reverse().get_custom() diff --git a/plain-postgres/tests/typing/select_ladder.py b/plain-postgres/tests/typing/select_ladder.py new file mode 100644 index 0000000000..99daaefeb3 --- /dev/null +++ b/plain-postgres/tests/typing/select_ladder.py @@ -0,0 +1,116 @@ +"""Every rung of `select()`'s overload ladder, column by column. + +The ladder is ten near-identical overloads, hand-written with no codegen, and +`select_rows.py` only exercised its ends -- one, two, three columns and the +cliff past ten. That left the middle rungs unasserted: transposing two +typevars in the six-column rung type-checked clean, so a wrong row type could +ship silently. + +Ten fields with ten *distinct* types is what makes that impossible. A +transposition inside any rung swaps two types in the asserted tuple, and the +`assert_type` fails. Nothing here runs -- the model is never registered at +runtime, because pytest doesn't collect this directory. +""" + +from __future__ import annotations + +from datetime import date, datetime, time, timedelta +from decimal import Decimal +from typing import assert_type +from uuid import UUID + +from plain.postgres import Field, RowQuerySet, types + +from plain import postgres + + +@postgres.register_model +class TenColumns(postgres.Model): + c0: Field[str] = types.TextField(max_length=10) + c1: Field[int] = types.IntegerField(default=0) + c2: Field[bool] = types.BooleanField(default=False) + c3: Field[float] = types.FloatField(default=0.0) + c4: Field[Decimal] = types.DecimalField(max_digits=5, decimal_places=2) + c5: Field[UUID] = types.UUIDField() + c6: Field[date] = types.DateField() + c7: Field[time] = types.TimeField() + c8: Field[datetime] = types.DateTimeField() + c9: Field[timedelta] = types.DurationField() + + +T = TenColumns + + +def must_accept_rung_one() -> None: + assert_type(T.query.select(T.c0), RowQuerySet[tuple[str]]) + + +def must_accept_rung_two() -> None: + assert_type(T.query.select(T.c0, T.c1), RowQuerySet[tuple[str, int]]) + + +def must_accept_rung_three() -> None: + assert_type(T.query.select(T.c0, T.c1, T.c2), RowQuerySet[tuple[str, int, bool]]) + + +def must_accept_rung_four() -> None: + assert_type( + T.query.select(T.c0, T.c1, T.c2, T.c3), + RowQuerySet[tuple[str, int, bool, float]], + ) + + +def must_accept_rung_five() -> None: + assert_type( + T.query.select(T.c0, T.c1, T.c2, T.c3, T.c4), + RowQuerySet[tuple[str, int, bool, float, Decimal]], + ) + + +def must_accept_rung_six() -> None: + assert_type( + T.query.select(T.c0, T.c1, T.c2, T.c3, T.c4, T.c5), + RowQuerySet[tuple[str, int, bool, float, Decimal, UUID]], + ) + + +def must_accept_rung_seven() -> None: + assert_type( + T.query.select(T.c0, T.c1, T.c2, T.c3, T.c4, T.c5, T.c6), + RowQuerySet[tuple[str, int, bool, float, Decimal, UUID, date]], + ) + + +def must_accept_rung_eight() -> None: + assert_type( + T.query.select(T.c0, T.c1, T.c2, T.c3, T.c4, T.c5, T.c6, T.c7), + RowQuerySet[tuple[str, int, bool, float, Decimal, UUID, date, time]], + ) + + +def must_accept_rung_nine() -> None: + assert_type( + T.query.select(T.c0, T.c1, T.c2, T.c3, T.c4, T.c5, T.c6, T.c7, T.c8), + RowQuerySet[tuple[str, int, bool, float, Decimal, UUID, date, time, datetime]], + ) + + +def must_accept_rung_ten() -> None: + assert_type( + T.query.select(T.c0, T.c1, T.c2, T.c3, T.c4, T.c5, T.c6, T.c7, T.c8, T.c9), + RowQuerySet[ + tuple[str, int, bool, float, Decimal, UUID, date, time, datetime, timedelta] + ], + ) + + +def must_accept_the_cliff_past_ten() -> None: + """An eleventh column has no rung; the *args fallback takes it.""" + from typing import Any + + assert_type( + T.query.select( + T.c0, T.c1, T.c2, T.c3, T.c4, T.c5, T.c6, T.c7, T.c8, T.c9, T.c0 + ), + RowQuerySet[tuple[Any, ...]], + ) diff --git a/plain-postgres/tests/typing/select_rows.py b/plain-postgres/tests/typing/select_rows.py new file mode 100644 index 0000000000..b2e23db6c0 --- /dev/null +++ b/plain-postgres/tests/typing/select_rows.py @@ -0,0 +1,250 @@ +"""`select()` resolves each column to a precise type in the row. + +`Field[T]` subclasses `Selectable[T]`, so a field contributes its `T` to the +row tuple; an expression subclasses `Selectable[Any]` and contributes `Any` +without blurring the fields beside it. The overload ladder on +`QuerySet.select()` is what binds those per-column typevars, and `flat=` and +`result_type=` each pick a different rung -- so the ladder is the promise, and +it is asserted statically. + +Nothing here runs, so the calls that would issue a real query (iterating, +first(), get()) are written out the same as any other claim. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Never, assert_type + +from app.examples.models.defaults import DefaultsExample as D +from app.examples.models.relationships import WidgetTag +from plain.postgres import Field, RowQuerySet, types +from plain.postgres.expressions import F +from plain.postgres.functions import Upper + +from plain import postgres + + +@dataclass +class NameStat: + name: str + priority: int + + +def must_accept_fields_as_per_column_types() -> None: + assert_type(D.query.select(D.name), RowQuerySet[tuple[str]]) + assert_type(D.query.select(D.name, D.priority), RowQuerySet[tuple[str, int]]) + assert_type( + D.query.select(D.name, D.priority, D.status), + RowQuerySet[tuple[str, int, str]], + ) + + +def must_accept_the_widest_rung_and_fall_back_past_it() -> None: + """The ladder stops at ten columns; an eleventh degrades to a plain tuple. + + Ten near-identical overloads are hand-written, so both the last rung and + the cliff past it are asserted -- a transposed typevar in a middle rung + would otherwise ship silently. + """ + assert_type( + D.query.select( + D.name, + D.priority, + D.status, + D.note, + D.id, + D.name, + D.priority, + D.status, + D.note, + D.id, + ), + RowQuerySet[ + tuple[str, int, str, str | None, int, str, int, str, str | None, int] + ], + ) + assert_type( + D.query.select( + D.name, + D.priority, + D.status, + D.note, + D.id, + D.name, + D.priority, + D.status, + D.note, + D.id, + D.name, + ), + RowQuerySet[tuple[Any, ...]], + ) + + +def must_accept_a_nullable_column_as_optional() -> None: + assert_type(D.query.select(D.name, D.note), RowQuerySet[tuple[str, str | None]]) + + +def must_accept_flat_as_the_bare_column_type() -> None: + assert_type(D.query.select(D.name, flat=True), RowQuerySet[str]) + assert_type(D.query.select(D.note, flat=True), RowQuerySet[str | None]) + + +def must_accept_result_type_as_the_row_type() -> None: + assert_type( + D.query.select(D.name, D.priority, result_type=NameStat), + RowQuerySet[NameStat], + ) + + +def must_accept_an_expression_as_any_without_blurring_its_neighbours() -> None: + # Expressions are Selectable[Any] for now, so only that column goes to Any. + assert_type(D.query.select(D.priority, Upper("name")), RowQuerySet[tuple[int, Any]]) + # F() is not a BaseExpression, but select() takes it like values_list does. + # Runtime half: tests/public/test_select.py::test_select_f_expression_column. + assert_type(D.query.select(D.name, F("priority")), RowQuerySet[tuple[str, Any]]) + assert_type(D.query.select(F("priority"), flat=True), RowQuerySet[Any]) + + +def must_accept_the_row_type_flowing_out_of_the_queryset() -> None: + for row in D.query.select(D.name, D.priority): + assert_type(row, tuple[str, int]) + assert_type(D.query.select(D.name, D.priority).first(), tuple[str, int] | None) + assert_type(D.query.select(D.name, D.priority).get(), tuple[str, int]) + assert_type( + D.query.select(D.name, D.priority).get_or_none(), tuple[str, int] | None + ) + for value in D.query.select(D.name, flat=True): + assert_type(value, str) + + +def must_reject_string_column_names() -> None: + # Runtime half: tests/public/test_select.py::test_select_rejects_string_argument. + D.query.select("name") # ty: ignore[no-matching-overload] + + +def must_reject_a_relation_reference() -> None: + # Runtime half: tests/public/test_select.py::test_select_rejects_fk_reference. + WidgetTag.query.select(WidgetTag.widget) # ty: ignore[no-matching-overload] + + +def must_reject_flat_with_more_than_one_column() -> None: + # flat= has a one-column rung only. Runtime half: test_select.py. + D.query.select(D.name, D.priority, flat=True) # ty: ignore[no-matching-overload] + + +def must_reject_flat_together_with_result_type() -> None: + # Runtime half: + # tests/public/test_select.py::test_select_flat_and_result_type_conflict. + D.query.select( # ty: ignore[no-matching-overload] + D.name, flat=True, result_type=NameStat + ) + + +def must_accept_row_mode_refusals_as_never_returning() -> None: + """The row-mode refusals are typed Never, so the checker knows they don't + return -- a call site's trailing code is unreachable rather than silently + typed as a QuerySet. + + `Never` on the *return* doesn't reject the call itself; the refusal is + still a runtime TypeError. Runtime half: + tests/public/test_select.py::test_values_after_select_raises and friends. + """ + rows = D.query.select(D.name) + assert_type(rows.values("name"), Never) + assert_type(rows.values_list("name"), Never) + assert_type(rows.get_or_create(name="a"), Never) + assert_type(rows.prefetch_related("tags"), Never) + assert_type(rows.bulk_update([], ["name"]), Never) + assert_type(rows.returning(), Never) + # annotate() appends a column, which would make the declared row type + # wrong. Runtime half: test_select.py::TestAnnotateAfterSelect. + assert_type(rows.annotate(n=Upper("name")), Never) + + +def must_accept_the_row_type_through_iterator_and_chaining() -> None: + rows = D.query.select(D.name, D.priority) + assert_type(rows.where(D.priority.gte(1)), RowQuerySet[tuple[str, int]]) + assert_type(rows[0], tuple[str, int]) + assert_type(rows[0:2], RowQuerySet[tuple[str, int]]) + for row in rows.iterator(): + assert_type(row, tuple[str, int]) + + +def must_accept_the_row_type_surviving_every_chaining_method() -> None: + """`RowQuerySet[R]` specializes its base as `QuerySet[Any]`, so an + inherited method annotated `QuerySet[T]` would hand back `QuerySet[Any]` + and drop `R`. They are annotated `Self` instead, which carries it. + """ + rows = D.query.select(D.name, D.priority) + assert_type(rows.reverse(), RowQuerySet[tuple[str, int]]) + assert_type(rows.none(), RowQuerySet[tuple[str, int]]) + assert_type(rows.distinct(), RowQuerySet[tuple[str, int]]) + assert_type(rows.order_by("name"), RowQuerySet[tuple[str, int]]) + assert_type(rows.for_update(), RowQuerySet[tuple[str, int]]) + assert_type(rows & rows, RowQuerySet[tuple[str, int]]) + assert_type(rows.all(), RowQuerySet[tuple[str, int]]) + + +def must_accept_or_keeping_the_row_type() -> None: + """`__or__` carries the row type too. + + It used to be the one chaining method that couldn't be `Self`: a sliced + left operand is re-expressed as an id subquery against + `Meta.base_queryset`, which is a plain `QuerySet` by design. #85 settled + it with an overload pair and a cast on that branch, so the row type + survives here like it does everywhere else. + """ + rows = D.query.select(D.name, D.priority) + assert_type(rows | rows, RowQuerySet[tuple[str, int]]) + + +def must_accept_a_column_from_another_model_because_the_checker_cannot_see_it() -> None: + """A column carries its value type but not its *model*. + + `Field[str]` is `Field[str]` whichever model declared it, so nothing here + distinguishes a column meant for `DefaultsExample.query.select()` from one + meant for `WidgetTag.query.select()`. This line has to type-check clean -- + that is the whole reason `select()` carries a runtime check instead, + mirroring `where()`'s in tests/typing/conditions_value_types.py. + + Runtime half: + tests/public/test_select.py::TestColumnsBelongToTheirModel. + """ + assert_type(D.query.select(WidgetTag.id), RowQuerySet[tuple[int]]) + + +@postgres.register_model +class ReadmeUser(postgres.Model): + """The README's lead `select()` example, so it can't rot. + + Mirrors `plain/postgres/README.md`'s "Selecting columns with select()" + block: the field annotations and the row type it claims have to keep + type-checking exactly as written there. + """ + + email: Field[str] = types.EmailField() + age: Field[int | None] = types.IntegerField(allow_null=True, default=None) + + +def must_accept_the_readme_example() -> None: + rows = ReadmeUser.query.where(ReadmeUser.age.gte(18)).select( + ReadmeUser.email, ReadmeUser.age + ) + assert_type(rows, RowQuerySet[tuple[str, int | None]]) + for email, age in rows: + assert_type(email, str) + assert_type(age, int | None) + + +def must_accept_selecting_after_returning_because_only_runtime_sees_it() -> None: + """The other order is a runtime-only refusal. + + `select()` after `returning()` raises, but nothing here can say so: + `ReturningQuerySet` is a `QuerySet` subclass, so `select()` resolves on it + like any other method. Only the `RowQuerySet` direction is typed (it + returns `Never`, asserted above). Runtime half: + tests/public/test_select.py::test_select_after_returning_raises. + """ + assert_type(D.query.returning().select(D.name), RowQuerySet[tuple[str]])