expressions

Author: mkurnikov

django-stubs/db/models/__init__.pyi

@@ -82,6 +82,12 @@ from .expressions import (
     Value as Value,
     Func as Func,
     ExpressionWrapper as ExpressionWrapper,
+    Combinable as Combinable,
+    Col as Col,
+    CombinedExpression as CombinedExpression,
+    ExpressionList as ExpressionList,
+    Random as Random,
+    Ref as Ref,
 )
 
 from .manager import BaseManager as BaseManager, Manager as Manager

django-stubs/db/models/expressions.pyi

@@ -1,19 +1,21 @@
 from collections import OrderedDict
 from datetime import datetime, timedelta
-from typing import Any, Callable, Dict, Iterator, List, Optional, Sequence, Set, Tuple, Type, Union
+from typing import Any, Callable, Dict, Iterator, List, Optional, Sequence, Set, Tuple, Type, Union, TypeVar
 
 from django.db.models.lookups import Lookup
 from django.db.models.sql.compiler import SQLCompiler
 
 from django.db.models import Q, QuerySet
-from django.db.models.fields import Field
+from django.db.models.fields import Field, FloatField
 from django.db.models.sql import Query
 
 _OutputField = Union[Field, str]
 
 class SQLiteNumericMixin:
     def as_sqlite(self, compiler: SQLCompiler, connection: Any, **extra_context: Any) -> Tuple[str, List[float]]: ...
 
+_SelfCombinable = TypeVar("_SelfCombinable", bound="Combinable")
+
 class Combinable:
     ADD: str = ...
     SUB: str = ...
@@ -25,28 +27,32 @@ class Combinable:
     BITOR: str = ...
     BITLEFTSHIFT: str = ...
     BITRIGHTSHIFT: str = ...
-    def __neg__(self) -> CombinedExpression: ...
-    def __add__(self, other: Optional[Union[timedelta, Combinable, float, str]]) -> CombinedExpression: ...
-    def __sub__(self, other: Union[timedelta, Combinable, float]) -> CombinedExpression: ...
-    def __mul__(self, other: Union[timedelta, Combinable, float]) -> CombinedExpression: ...
-    def __truediv__(self, other: Union[Combinable, float]) -> Combinable: ...
-    def __itruediv__(self, other: Union[Combinable, float]) -> Combinable: ...
-    def __mod__(self, other: int) -> Combinable: ...
-    def __pow__(self, other: float) -> Combinable: ...
-    def __and__(self, other: Combinable) -> Any: ...
-    def bitand(self, other: int) -> CombinedExpression: ...
-    def bitleftshift(self, other: int) -> CombinedExpression: ...
-    def bitrightshift(self, other: int) -> CombinedExpression: ...
-    def __or__(self, other: Combinable) -> Any: ...
-    def bitor(self, other: int) -> CombinedExpression: ...
-    def __radd__(self, other: Optional[Union[datetime, float]]) -> CombinedExpression: ...
-    def __rsub__(self, other: float) -> CombinedExpression: ...
-    def __rmul__(self, other: float) -> CombinedExpression: ...
-    def __rtruediv__(self, other: float) -> CombinedExpression: ...
-    def __rmod__(self, other: int) -> CombinedExpression: ...
-    def __rpow__(self, other: float) -> CombinedExpression: ...
-    def __rand__(self, other: Any) -> Any: ...
-    def __ror__(self, other: Any) -> Any: ...
+    def __neg__(self: _SelfCombinable) -> _SelfCombinable: ...
+    def __add__(
+        self: _SelfCombinable, other: Optional[Union[timedelta, Combinable, float, str]]
+    ) -> _SelfCombinable: ...
+    def __sub__(self: _SelfCombinable, other: Union[timedelta, Combinable, float]) -> _SelfCombinable: ...
+    def __mul__(self: _SelfCombinable, other: Union[timedelta, Combinable, float]) -> _SelfCombinable: ...
+    def __truediv__(self: _SelfCombinable, other: Union[Combinable, float]) -> _SelfCombinable: ...
+    def __itruediv__(self: _SelfCombinable, other: Union[Combinable, float]) -> _SelfCombinable: ...
+    def __mod__(self: _SelfCombinable, other: Union[int, Combinable]) -> _SelfCombinable: ...
+    def __pow__(self: _SelfCombinable, other: Union[float, Combinable]) -> _SelfCombinable: ...
+    def __and__(self: _SelfCombinable, other: Combinable) -> _SelfCombinable: ...
+    def bitand(self: _SelfCombinable, other: int) -> _SelfCombinable: ...
+    def bitleftshift(self: _SelfCombinable, other: int) -> _SelfCombinable: ...
+    def bitrightshift(self: _SelfCombinable, other: int) -> _SelfCombinable: ...
+    def __or__(self: _SelfCombinable, other: Combinable) -> _SelfCombinable: ...
+    def bitor(self: _SelfCombinable, other: int) -> _SelfCombinable: ...
+    def __radd__(self, other: Optional[Union[datetime, float, Combinable]]) -> Combinable: ...
+    def __rsub__(self, other: Union[float, Combinable]) -> Combinable: ...
+    def __rmul__(self, other: Union[float, Combinable]) -> Combinable: ...
+    def __rtruediv__(self, other: Union[float, Combinable]) -> Combinable: ...
+    def __rmod__(self, other: Union[int, Combinable]) -> Combinable: ...
+    def __rpow__(self, other: Union[float, Combinable]) -> Combinable: ...
+    def __rand__(self, other: Any) -> Combinable: ...
+    def __ror__(self, other: Any) -> Combinable: ...
+
+_SelfBaseExpression = TypeVar("_SelfBaseExpression", bound="BaseExpression")
 
 class BaseExpression:
     is_summary: bool = ...
@@ -74,9 +80,9 @@ class BaseExpression:
     def convert_value(self) -> Callable: ...
     def get_lookup(self, lookup: str) -> Optional[Type[Lookup]]: ...
     def get_transform(self, name: str) -> Optional[Type[Expression]]: ...
-    def relabeled_clone(self, change_map: Union[Dict[Optional[str], str], OrderedDict]) -> Expression: ...
+    def relabeled_clone(self, change_map: Dict[Optional[str], str]) -> Expression: ...
     def copy(self) -> BaseExpression: ...
-    def get_group_by_cols(self) -> List[Expression]: ...
+    def get_group_by_cols(self: _SelfBaseExpression) -> List[_SelfBaseExpression]: ...
     def get_source_fields(self) -> List[Optional[Field]]: ...
     def asc(self, **kwargs: Any) -> Expression: ...
     def desc(self, **kwargs: Any) -> Expression: ...
@@ -125,7 +131,8 @@ class F(Combinable):
     def desc(self, **kwargs) -> OrderBy: ...
     def deconstruct(self) -> Any: ...
 
-class OuterRef(F): ...
+class OuterRef(F):
+    def __init__(self, name: Union[str, OuterRef]): ...
 
 class Subquery(Expression):
     template: str = ...
@@ -197,6 +204,18 @@ class Case(Expression):
     ) -> None: ...
 
 class ExpressionWrapper(Expression):
-    def __init__(self, expression: Union[Expression, Q], output_field: _OutputField): ...
+    def __init__(self, expression: Union[Q, Combinable], output_field: _OutputField): ...
     def set_source_expressions(self, exprs: Sequence[Expression]) -> None: ...
     def get_source_expressions(self) -> List[Expression]: ...
+
+class Col(Expression):
+    def __init__(self, alias: str, target: str, output_field: Optional[_OutputField] = ...): ...
+
+class ExpressionList(Func):
+    def __init__(self, *expressions: Union[BaseExpression, Combinable], **extra: Any) -> None: ...
+
+class Random(Expression):
+    output_field: FloatField
+
+class Ref(Expression):
+    def __init__(self, refs: str, source: Expression): ...

django-stubs/db/models/fields/__init__.pyi

@@ -2,6 +2,8 @@ from typing import Any, Optional, Tuple, Iterable, Callable, Dict, Union, Type
 
 from django.db.models import Model
 from django.db.models.query_utils import RegisterLookupMixin
+
+from django.db.models.expressions import F, Combinable
 from django.forms import Widget, Field as FormField
 
 _Choice = Tuple[Any, str]
@@ -52,6 +54,7 @@ class Field(RegisterLookupMixin):
     def contribute_to_class(self, cls: Type[Model], name: str, private_only: bool = ...) -> None: ...
 
 class IntegerField(Field):
+    def __set__(self, instance, value: Union[int, F]) -> None: ...
     def __get__(self, instance, owner) -> int: ...
 
 class PositiveIntegerRelDbTypeMixin:
@@ -115,7 +118,7 @@ class CharField(Field):
         validators: Iterable[_ValidatorCallable] = ...,
         error_messages: Optional[_ErrorMessagesToOverride] = ...,
     ): ...
-    def __set__(self, instance, value: str) -> None: ...
+    def __set__(self, instance, value: Union[str, Combinable]) -> None: ...
     def __get__(self, instance, owner) -> str: ...
 
 class SlugField(CharField):

django-stubs/db/models/fields/related.pyi

@@ -32,6 +32,8 @@ from django.db.models.fields.reverse_related import (
 )
 from django.db.models.query_utils import PathInfo, Q
 
+from django.db.models.expressions import F
+
 if TYPE_CHECKING:
     from django.db.models.manager import RelatedManager
 
@@ -102,6 +104,7 @@ class ForeignObject(RelatedField):
 
 class ForeignKey(RelatedField, Generic[_T]):
     def __init__(self, to: Union[Type[_T], str], on_delete: Any, related_name: str = ..., **kwargs): ...
+    def __set__(self, instance, value: Union[Model, F]) -> None: ...
     def __get__(self, instance, owner) -> _T: ...
 
 class OneToOneField(RelatedField, Generic[_T]):

scripts/typecheck_tests.py

@@ -201,7 +201,7 @@
     # TODO: 'dispatch',
     'distinct_on_fields',
     'empty',
-    # TODO: 'expressions',
+    'expressions',
     'expressions_case',
     # TODO: 'expressions_window',
     # TODO: 'extra_regress',