Update typings to be more precise (#42)

Author: bryanforbes

sqlalchemy-stubs/sql/elements.pyi

@@ -13,6 +13,7 @@ from .type_api import TypeEngine
 from .sqltypes import NullType, Boolean, Integer
 from .selectable import TextAsFrom, TableClause
 from .functions import FunctionElement
+from .schema import ForeignKey
 
 _T = TypeVar('_T')
 _T_contra = TypeVar('_T_contra', contravariant=True)
@@ -44,27 +45,32 @@ class ClauseElement(Visitable):
     def __bool__(self): ...
     __nonzero__: Any = ...
 
+_CE = TypeVar('_CE', bound=ColumnElement)
+
 class ColumnElement(operators.ColumnOperators, ClauseElement, Generic[_T]):
     __visit_name__: str = ...
-    primary_key: Any = ...
-    foreign_keys: List[Any] = ...
-    key: Any = ...
-    def self_group(self, against: Optional[Any] = ...): ...
+    primary_key: bool = ...
+    foreign_keys: List[ForeignKey] = ...
+    key: Optional[str] = ...
+    def self_group(self: _CE, against: Optional[Any] = ...) -> Union[AsBoolean, Grouping[_T], _CE]: ...
     @property
     def type(self) -> TypeEngine[_T]: ...
     def comparator(self): ...
-    def __getattr__(self, key): ...
+    def __getattr__(self, key: str) -> Any: ...
     def operate(self, op, *other, **kwargs): ...
     def reverse_operate(self, op, other, **kwargs): ...
     @property
-    def expression(self): ...
-    def base_columns(self): ...
-    def proxy_set(self): ...
-    def shares_lineage(self, othercolumn): ...
-    def compare(self, other: Any, **kw: Any) -> bool: ...
-    def cast(self, type_): ...
-    def label(self, name): ...
-    def anon_label(self): ...
+    def expression(self: _CE) -> _CE: ...
+    def base_columns(self) -> Set[ColumnElement[Any]]: ...
+    def proxy_set(self) -> Set[ColumnElement[Any]]: ...
+    def shares_lineage(self, othercolumn: ColumnElement[Any]) -> bool: ...
+    def compare(self, other: ColumnElement[Any], use_proxies: bool = ..., equivalents: bool = ..., **kw: Any) -> bool: ...
+    @overload
+    def cast(self, type_: Type[TypeEngine[_U]]) -> Cast[_U]: ...
+    @overload
+    def cast(self, type_: TypeEngine[_U]) -> Cast[_U]: ...
+    def label(self, name: str) -> Label[_T]: ...
+    def anon_label(self) -> _anonymous_label: ...
 
 class BindParameter(ColumnElement[_T]):
     __visit_name__: str = ...
@@ -87,7 +93,8 @@ class BindParameter(ColumnElement[_T]):
                  _compared_to_type: Optional[Any] = ...) -> None: ...
     @property
     def effective_value(self) -> _T: ...
-    def compare(self, other: Any, **kw: Any) -> bool: ...
+    # Signature of "compare" incompatible with supertype "ColumnElement"
+    def compare(self, other: ColumnElement[Any], **kw: Any) -> bool: ...  # type: ignore
 
 class TypeClause(ClauseElement, Generic[_T]):
     __visit_name__: str = ...
@@ -168,19 +175,21 @@ class ClauseList(ClauseElement):
 
 _BCL = TypeVar('_BCL', bound=BooleanClauseList)
 
-class BooleanClauseList(ClauseList, ColumnElement[bool]):
+# Definitions of "compare" and "self_group" in ClauseList are incompatible with ColumnElement
+class BooleanClauseList(ClauseList, ColumnElement[bool]):  # type: ignore
     __visit_name__: str = ...
     # Note: passing strings to below generates a warning, but still works.
     @classmethod
     def and_(cls, *clauses: Union[ClauseElement, str, bool]) -> BooleanClauseList: ...
     @classmethod
     def or_(cls, *clauses: Union[ClauseElement, str, bool]) -> BooleanClauseList: ...
-    def self_group(self: _BCL, against: Optional[Any] = ...) -> Union[_BCL, Grouping[Any]]: ...
+    def self_group(self: _BCL, against: Optional[Any] = ...) -> Union[_BCL, Grouping[bool]]: ...
 
 and_ = BooleanClauseList.and_
 or_ = BooleanClauseList.or_
 
-class Tuple(ClauseList, ColumnElement[_T]):
+# Definitions of "compare" and "self_group" in ClauseList are incompatible with ColumnElement
+class Tuple(ClauseList, ColumnElement[_T]):  # type: ignore
     type: TypeEngine[_T] = ...
     @overload
     def __init__(self, *clauses: ColumnElement[Any], type_: Type[TypeEngine[_T]], **kw: Any) -> None: ...
@@ -261,7 +270,8 @@ class UnaryExpression(ColumnElement[_T]):
                  type_: Optional[TypeEngine[_T]] = ..., negate: Optional[Any] = ...,
                  wraps_column_expression: bool = ...) -> None: ...
     def get_children(self, **kwargs) -> _TupleType[Any]: ...
-    def compare(self, other: Any, **kw: Any) -> bool: ...
+    # Signature of "compare" incompatible with supertype "ColumnElement"
+    def compare(self, other: ColumnElement[Any], **kw: Any) -> bool: ...  # type: ignore
     def self_group(self: _UE, against: Optional[Any] = ...) -> Union[_UE, Grouping[_T]]: ...
     @classmethod
     def _create_nullsfirst(cls, column: ColumnElement[Any]) -> UnaryExpression[NullType]: ...
@@ -292,14 +302,14 @@ class CollectionAggregate(UnaryExpression[_T]):
 
 _AB = TypeVar('_AB', bound=AsBoolean)
 
-class AsBoolean(UnaryExpression):
-    element: Any = ...
+class AsBoolean(UnaryExpression[bool]):
+    element: ColumnElement[Any] = ...
     type: Boolean = ...
     operator: Callable[..., Any] = ...
-    negate: Any = ...
+    negate: Callable[..., Any] = ...
     modifier: Any = ...
     wraps_column_expression: bool = ...
-    def __init__(self, element: Any, operator: Callable[..., Any], negate: Any) -> None: ...
+    def __init__(self, element: ColumnElement[Any], operator: Callable[..., Any], negate: Callable[..., Any]) -> None: ...
     def self_group(self: _AB, against: Optional[Any] = ...) -> _AB: ...
 
 _BE = TypeVar('_BE', bound=BinaryExpression)
@@ -324,7 +334,8 @@ class BinaryExpression(ColumnElement[_T], Generic[_T, _U, _V]):
     @property
     def is_comparison(self): ...
     def get_children(self, **kwargs): ...
-    def compare(self, other: Any, **kw: Any) -> bool: ...
+    # Signature of "compare" incompatible with supertype "ColumnElement"
+    def compare(self, other: ColumnElement[Any], **kw: Any) -> bool: ...  # type: ignore
     def self_group(self: _BE, against: Optional[Any] = ...) -> Union[_BE, Grouping[_T]]: ...
 
 _SL = TypeVar('_SL', bound=Slice)
@@ -336,7 +347,7 @@ class Slice(ColumnElement[None]):
     step: Any = ...
     type: NullType = ...
     def __init__(self, start, stop, step) -> None: ...
-    def self_group(self: _SL, against: Optional[Any] = ...) -> Union[_SL, Grouping[NullType]]: ...
+    def self_group(self: _SL, against: Optional[Any] = ...) -> _SL: ...
 
 class IndexExpression(BinaryExpression): ...
 
@@ -354,7 +365,8 @@ class Grouping(ColumnElement[_T]):
     def self_group(self: _G, against: Optional[Any] = ...) -> _G: ...
     def get_children(self, **kwargs) -> _TupleType[ClauseElement]: ...
     def __getattr__(self, attr): ...
-    def compare(self, other: Any, **kw: Any) -> bool: ...
+    # Signature of "compare" incompatible with supertype "ColumnElement"
+    def compare(self, other: ColumnElement[Any], **kw: Any) -> bool: ...  # type: ignore
 
 RANGE_UNBOUNDED: util.symbol = ...
 RANGE_CURRENT: util.symbol = ...
@@ -432,11 +444,13 @@ class Label(ColumnElement[_T]):
     def type(self) -> TypeEngine[_T]: ...
     @property
     def element(self) -> ColumnElement[_T]: ...
-    def self_group(self, against: Optional[Any] = ...) -> Union[ColumnElement[_T], Grouping[_T]]: ...
+    def self_group(self: _L, against: Optional[Any] = ...) -> Union[_L, Label[_T]]: ...
+    # Signature of "primary_key" incompatible with supertype "ColumnElement"
     @property
-    def primary_key(self) -> Any: ...
+    def primary_key(self) -> bool: ...  # type: ignore
+    # Signature of "foreign_keys" incompatible with supertype "ColumnElement"
     @property
-    def foreign_keys(self) -> List[Any]: ...  # type: ignore
+    def foreign_keys(self) -> List[ForeignKey]: ...  # type: ignore
     def get_children(self, **kwargs: Any) -> Any: ...
 
 class ColumnClause(Immutable, ColumnElement[_T]):
@@ -485,6 +499,8 @@ class _truncated_label(quoted_name):
 
 class conv(_truncated_label): ...
 
+class _anonymous_label(_truncated_label): ...
+
 class AnnotatedColumnElement(Annotated, Generic[_T]):
     def __init__(self, element: ColumnElement[_T], values: Any) -> None: ...
     @property

sqlalchemy-stubs/sql/functions.pyi

@@ -1,6 +1,8 @@
 from typing import Any, Optional, Union, Iterable, Tuple, TypeVar, Generic, overload
 from .base import Executable as Executable, ColumnCollection
-from .elements import ColumnElement as ColumnElement, Grouping, ClauseList, ColumnElement, Over, WithinGroup, FunctionFilter
+from .elements import (
+    ColumnElement as ColumnElement, Grouping, ClauseList, ColumnElement, Over, WithinGroup, FunctionFilter, AsBoolean
+)
 from .selectable import FromClause as FromClause, Alias, Select
 from . import util as sqlutil
 from .visitors import VisitableType as VisitableType
@@ -13,6 +15,8 @@ _T = TypeVar('_T')
 
 def register_function(identifier, fn, package: str = ...): ...
 
+_FE = TypeVar('_FE', bound=FunctionElement)
+
 class FunctionElement(Executable, ColumnElement[_T], FromClause, Generic[_T]):  # type: ignore
     # ColumnElement.foreign_keys() is not compatible with FromClause.foreign_keys()
     packagenames: Any = ...
@@ -35,7 +39,7 @@ class FunctionElement(Executable, ColumnElement[_T], FromClause, Generic[_T]):
                                             ColumnElement[Any],
                                             Iterable[Union[str, ColumnElement[Any]]]]) -> WithinGroup[_T]: ...
     @overload
-    def filter(self) -> FunctionElement[_T]: ...
+    def filter(self: _FE) -> _FE: ...
     @overload
     def filter(self, criteria: Any, *criterion: Any) -> FunctionFilter[_T]: ...
     def get_children(self, **kwargs: Any) -> Tuple[Grouping[Any]]: ...
@@ -44,7 +48,7 @@ class FunctionElement(Executable, ColumnElement[_T], FromClause, Generic[_T]):
     def select(self) -> Select: ...  # type: ignore  # incompatible with FromClause.select
     def scalar(self) -> Any: ...  # type: ignore  # incompatible with Executable.scalar
     def execute(self) -> ResultProxy: ...  # type: ignore  # incompatible with Executable.execute
-    def self_group(self, against: Optional[Any] = ...) -> Union[ColumnElement[_T], Grouping[_T]]: ...
+    def self_group(self: _FE, against: Optional[Any] = ...) -> Union[AsBoolean, Grouping[_T], _FE]: ...
 
 class _FunctionGenerator(object):
     opts: Any = ...