update import_all test, add some stubs for postgres (#178)

django-stubs/contrib/postgres/aggregates/__init__.pyi

@@ -0,0 +1,24 @@
+from .general import (
+    ArrayAgg as ArrayAgg,
+    BitAnd as BitAnd,
+    BitOr as BitOr,
+    BoolAnd as BoolAnd,
+    BoolOr as BoolOr,
+    JSONBAgg as JSONBAgg,
+    StringAgg as StringAgg,
+)
+
+from .statistics import (
+    Corr as Corr,
+    CovarPop as CovarPop,
+    RegrAvgX as RegrAvgX,
+    RegrAvgY as RegrAvgY,
+    RegrCount as RegrCount,
+    RegrIntercept as RegrIntercept,
+    RegrR2 as RegrR2,
+    RegrSlope as RegrSlope,
+    RegrSXX as RegrSXX,
+    RegrSXY as RegrSXY,
+    RegrSYY as RegrSYY,
+    StatAggregate as StatAggregate,
+)

django-stubs/contrib/postgres/aggregates/general.pyi

@@ -0,0 +1,11 @@
+from django.db.models.aggregates import Aggregate
+
+from .mixins import OrderableAggMixin
+
+class ArrayAgg(OrderableAggMixin, Aggregate): ...
+class BitAnd(Aggregate): ...
+class BitOr(Aggregate): ...
+class BoolAnd(Aggregate): ...
+class BoolOr(Aggregate): ...
+class JSONBAgg(Aggregate): ...
+class StringAgg(OrderableAggMixin, Aggregate): ...

django-stubs/contrib/postgres/aggregates/mixins.pyi

@@ -0,0 +1 @@
+class OrderableAggMixin: ...

django-stubs/contrib/postgres/aggregates/statistics.pyi

@@ -0,0 +1,14 @@
+from django.db.models.aggregates import Aggregate
+
+class StatAggregate(Aggregate): ...
+class Corr(StatAggregate): ...
+class CovarPop(StatAggregate): ...
+class RegrAvgX(StatAggregate): ...
+class RegrAvgY(StatAggregate): ...
+class RegrCount(StatAggregate): ...
+class RegrIntercept(StatAggregate): ...
+class RegrR2(StatAggregate): ...
+class RegrSlope(StatAggregate): ...
+class RegrSXX(StatAggregate): ...
+class RegrSXY(StatAggregate): ...
+class RegrSYY(StatAggregate): ...

django-stubs/contrib/postgres/indexes.pyi

@@ -1,29 +1,81 @@
-from typing import Any, Optional
+from typing import Any, Optional, Sequence
 
 from django.db.models import Index
+from django.db.models.query_utils import Q
 
-class PostgresIndex(Index):
-    @property
-    def max_name_length(self) -> int: ...
+class PostgresIndex(Index): ...
 
 class BrinIndex(PostgresIndex):
     def __init__(
-        self, *, autosummarize: Optional[bool] = ..., pages_per_range: Optional[int] = ..., **kwargs: Any
+        self,
+        *,
+        autosummarize: Optional[bool] = ...,
+        pages_per_range: Optional[int] = ...,
+        fields: Sequence[str] = ...,
+        name: Optional[str] = ...,
+        db_tablespace: Optional[str] = ...,
+        opclasses: Sequence[str] = ...,
+        condition: Optional[Q] = ...
     ) -> None: ...
 
 class BTreeIndex(PostgresIndex):
-    def __init__(self, *, fillfactor: Optional[int] = ..., **kwargs: Any): ...
+    def __init__(
+        self,
+        *,
+        fillfactor: Optional[int] = ...,
+        fields: Sequence[str] = ...,
+        name: Optional[str] = ...,
+        db_tablespace: Optional[str] = ...,
+        opclasses: Sequence[str] = ...,
+        condition: Optional[Q] = ...
+    ) -> None: ...
 
 class GinIndex(PostgresIndex):
     def __init__(
-        self, *, fastupdate: Optional[bool] = ..., gin_pending_list_limit: Optional[int] = ..., **kwargs: Any
+        self,
+        *,
+        fastupdate: Optional[bool] = ...,
+        gin_pending_list_limit: Optional[int] = ...,
+        fields: Sequence[str] = ...,
+        name: Optional[str] = ...,
+        db_tablespace: Optional[str] = ...,
+        opclasses: Sequence[str] = ...,
+        condition: Optional[Q] = ...
     ) -> None: ...
 
 class GistIndex(PostgresIndex):
-    def __init__(self, *, buffering: Optional[bool] = ..., fillfactor: Optional[int] = ..., **kwargs: Any) -> None: ...
+    def __init__(
+        self,
+        *,
+        buffering: Optional[bool] = ...,
+        fillfactor: Optional[int] = ...,
+        fields: Sequence[str] = ...,
+        name: Optional[str] = ...,
+        db_tablespace: Optional[str] = ...,
+        opclasses: Sequence[str] = ...,
+        condition: Optional[Q] = ...
+    ) -> None: ...
 
 class HashIndex(PostgresIndex):
-    def __init__(self, *, fillfactor: Optional[int] = ..., **kwargs: Any) -> None: ...
+    def __init__(
+        self,
+        *,
+        fillfactor: Optional[int] = ...,
+        fields: Sequence[str] = ...,
+        name: Optional[str] = ...,
+        db_tablespace: Optional[str] = ...,
+        opclasses: Sequence[str] = ...,
+        condition: Optional[Q] = ...
+    ) -> None: ...
 
 class SpGistIndex(PostgresIndex):
-    def __init__(self, *, fillfactor: Optional[int] = ..., **kwargs: Any) -> None: ...
+    def __init__(
+        self,
+        *,
+        fillfactor: Optional[int] = ...,
+        fields: Sequence[str] = ...,
+        name: Optional[str] = ...,
+        db_tablespace: Optional[str] = ...,
+        opclasses: Sequence[str] = ...,
+        condition: Optional[Q] = ...
+    ) -> None: ...

django-stubs/contrib/postgres/lookups.pyi

@@ -0,0 +1,18 @@
+from django.db.models.lookups import Exact
+
+from django.db.models import Lookup, Transform
+from .search import SearchVectorExact
+
+class PostgresSimpleLookup(Lookup):
+    operator: str
+
+class DataContains(PostgresSimpleLookup): ...
+class ContainedBy(PostgresSimpleLookup): ...
+class Overlap(PostgresSimpleLookup): ...
+class HasKey(PostgresSimpleLookup): ...
+class HasKeys(PostgresSimpleLookup): ...
+class HasAnyKeys(HasKeys): ...
+class Unaccent(Transform): ...
+class SearchLookup(SearchVectorExact): ...
+class TrigramSimilar(PostgresSimpleLookup): ...
+class JSONExact(Exact): ...

django-stubs/contrib/postgres/search.pyi

@@ -1,10 +1,12 @@
 from typing import Any, Dict, Optional, TypeVar, Union
 
-from django.db.models.expressions import Combinable, CombinedExpression, Func, Value
+from django.db.models.expressions import Combinable, CombinedExpression, Func, Value, _OutputField
 from django.db.models.lookups import Lookup
 
 from django.db.models import Field
 
+_Expression = Union[str, Combinable, "SearchQueryCombinable"]
+
 class SearchVectorExact(Lookup): ...
 class SearchVectorField(Field): ...
 class SearchQueryField(Field): ...
@@ -14,10 +16,12 @@ class SearchVectorCombinable:
 
 class SearchVector(SearchVectorCombinable, Func):
     config: Optional[Any] = ...
-    def __init__(self, *expressions: Union[str, Combinable], **extra: Any): ...
+    def __init__(self, *expressions: _Expression, **extra: Any): ...
 
 class CombinedSearchVector(SearchVectorCombinable, CombinedExpression):
-    def __init__(self, lhs, connector, rhs, config, output_field: Optional[Field, str] = ...): ...
+    def __init__(
+        self, lhs, connector, rhs, config: Optional[_Expression] = ..., output_field: Optional[_OutputField] = ...
+    ): ...
 
 _T = TypeVar("_T", bound="SearchQueryCombinable")
 
@@ -29,19 +33,31 @@ class SearchQueryCombinable:
     def __and__(self: _T, other: SearchQueryCombinable) -> _T: ...
     def __rand__(self: _T, other: SearchQueryCombinable) -> _T: ...
 
-class SearchQuery(SearchQueryCombinable, Value):
-    SEARCH_TYPES: Dict[str, str] = {"plain": "plainto_tsquery", "phrase": "phraseto_tsquery", "raw": "to_tsquery"}
-    def __init__(self, value, output_field=..., *, config=..., invert=False, search_type="plain"): ...
+class SearchQuery(SearchQueryCombinable, Value):  # type: ignore
+    SEARCH_TYPES: Dict[str, str] = ...
+    def __init__(
+        self,
+        value: str,
+        output_field: Optional[_OutputField] = ...,
+        *,
+        config: Optional[_Expression] = ...,
+        invert: bool = ...,
+        search_type: str = ...
+    ): ...
     def __invert__(self: _T) -> _T: ...
 
-class CombinedSearchQuery(SearchQueryCombinable, CombinedExpression):
-    def __init__(self, lhs, connector, rhs, config, output_field=...) -> None: ...
+class CombinedSearchQuery(SearchQueryCombinable, CombinedExpression):  # type: ignore
+    def __init__(
+        self, lhs, connector, rhs, config: Optional[_Expression] = ..., output_field: Optional[_OutputField] = ...
+    ) -> None: ...
 
 class SearchRank(Func):
-    def __init__(self, vector, query, **extra: Any) -> None: ...
+    def __init__(
+        self, vector: Union[SearchVector, _Expression], query: Union[SearchQuery, _Expression], **extra: Any
+    ) -> None: ...
 
 class TrigramBase(Func):
-    def __init__(self, expression, string, **extra: Any) -> None: ...
+    def __init__(self, expression: _Expression, string, **extra: Any) -> None: ...
 
 class TrigramSimilarity(TrigramBase): ...
 class TrigramDistance(TrigramBase): ...

django-stubs/contrib/staticfiles/testing.pyi

@@ -0,0 +1,3 @@
+from django.test import LiveServerTestCase
+
+class StaticLiveServerTestCase(LiveServerTestCase): ...

django-stubs/db/models/aggregates.pyi

@@ -5,6 +5,7 @@ from django.db.models.expressions import Func
 class Aggregate(Func):
     filter_template: str = ...
     filter: Any = ...
+    allow_distinct: bool = ...
     def __init__(self, *expressions: Any, distinct: bool = ..., filter: Optional[Any] = ..., **extra: Any) -> None: ...
 
 class Avg(Aggregate): ...