SQLAlchemy improvements for Operators (#7604)

stubs/SQLAlchemy/sqlalchemy/ext/associationproxy.pyi

@@ -77,7 +77,7 @@ class ObjectAssociationProxyInstance(AssociationProxyInstance):
     def __ne__(self, obj): ...
 
 class ColumnAssociationProxyInstance(ColumnOperators[Any], AssociationProxyInstance):
-    def __eq__(self, other): ...
+    def __eq__(self, other) -> ColumnOperators[Any]: ...  # type: ignore[override]
     def operate(self, op, *other, **kwargs): ...
 
 class _lazy_collection:

stubs/SQLAlchemy/sqlalchemy/ext/hybrid.pyi

@@ -44,7 +44,7 @@ class hybrid_property(interfaces.InspectionAttrInfo):
     def comparator(self, comparator): ...
     def update_expression(self, meth): ...
 
-class Comparator(interfaces.PropComparator):
+class Comparator(interfaces.PropComparator[Any]):
     property: Any
     expression: Any
     def __init__(self, expression) -> None: ...

stubs/SQLAlchemy/sqlalchemy/orm/attributes.pyi

@@ -13,7 +13,7 @@ NO_KEY: Any
 class QueryableAttribute(
     interfaces._MappedAttribute,
     interfaces.InspectionAttr,
-    interfaces.PropComparator,
+    interfaces.PropComparator[Any],
     traversals.HasCopyInternals,
     roles.JoinTargetRole,
     roles.OnClauseRole,

stubs/SQLAlchemy/sqlalchemy/orm/descriptor_props.pyi

@@ -1,9 +1,12 @@
-from typing import Any
+from typing import Any, Generic, TypeVar
 
+from ..sql.operators import ColumnOperators
 from ..util import memoized_property
 from . import util as orm_util
 from .interfaces import MapperProperty, PropComparator
 
+_T = TypeVar("_T")
+
 class DescriptorProperty(MapperProperty):
     doc: Any
     uses_objects: bool
@@ -33,15 +36,15 @@ class CompositeProperty(DescriptorProperty):
         def __init__(self, property_, expr) -> None: ...
         def create_row_processor(self, query, procs, labels): ...
 
-    class Comparator(PropComparator):
+    class Comparator(PropComparator[_T], Generic[_T]):
         __hash__: Any
         @memoized_property
         def clauses(self): ...
         def __clause_element__(self): ...
         @memoized_property
         def expression(self): ...
-        def __eq__(self, other): ...
-        def __ne__(self, other): ...
+        def __eq__(self, other) -> ColumnOperators[_T]: ...  # type: ignore[override]
+        def __ne__(self, other) -> ColumnOperators[_T]: ...  # type: ignore[override]
 
 class ConcreteInheritedProperty(DescriptorProperty):
     descriptor: Any

stubs/SQLAlchemy/sqlalchemy/orm/interfaces.pyi

@@ -1,4 +1,4 @@
-from typing import Any
+from typing import Any, Generic, TypeVar
 
 from .. import util
 from ..sql import operators, roles
@@ -17,6 +17,8 @@ from .base import (
     _MappedAttribute as _MappedAttribute,
 )
 
+_T = TypeVar("_T")
+
 __all__ = (
     "EXT_CONTINUE",
     "EXT_STOP",
@@ -57,7 +59,7 @@ class MapperProperty(HasCacheKey, _MappedAttribute, InspectionAttr, util.Memoize
         self, session, source_state, source_dict, dest_state, dest_dict, load, _recursive, _resolve_conflict_map
     ) -> None: ...
 
-class PropComparator(operators.ColumnOperators):
+class PropComparator(operators.ColumnOperators[_T], Generic[_T]):
     __visit_name__: str
     prop: Any
     property: Any

stubs/SQLAlchemy/sqlalchemy/orm/properties.pyi

@@ -38,7 +38,7 @@ class ColumnProperty(StrategizedProperty):
         self, session, source_state, source_dict, dest_state, dest_dict, load, _recursive, _resolve_conflict_map
     ) -> None: ...
 
-    class Comparator(util.MemoizedSlots, PropComparator):
+    class Comparator(util.MemoizedSlots, PropComparator[Any]):
         expressions: Any
         def _memoized_method___clause_element__(self): ...
         def operate(self, op, *other, **kwargs): ...

stubs/SQLAlchemy/sqlalchemy/orm/relationships.pyi

@@ -1,8 +1,11 @@
-from typing import Any
+from typing import Any, Generic, TypeVar
 
+from ..sql.operators import ColumnOperators
 from ..util import memoized_property
 from .interfaces import PropComparator, StrategizedProperty
 
+_T = TypeVar("_T")
+
 def remote(expr): ...
 def foreign(expr): ...
 
@@ -84,7 +87,7 @@ class RelationshipProperty(StrategizedProperty):
     ) -> None: ...
     def instrument_class(self, mapper) -> None: ...
 
-    class Comparator(PropComparator):
+    class Comparator(PropComparator[_T], Generic[_T]):
         prop: Any
         def __init__(
             self, prop, parentmapper, adapt_to_entity: Any | None = ..., of_type: Any | None = ..., extra_criteria=...
@@ -97,13 +100,13 @@ class RelationshipProperty(StrategizedProperty):
         def __clause_element__(self): ...
         def of_type(self, cls): ...
         def and_(self, *other): ...
-        def in_(self, other) -> None: ...
+        def in_(self, other) -> ColumnOperators[_T]: ...
         __hash__: Any
         def __eq__(self, other): ...
         def any(self, criterion: Any | None = ..., **kwargs): ...
         def has(self, criterion: Any | None = ..., **kwargs): ...
-        def contains(self, other, **kwargs): ...
-        def __ne__(self, other): ...
+        def contains(self, other, **kwargs) -> ColumnOperators[_T]: ...
+        def __ne__(self, other) -> ColumnOperators[_T]: ...  # type: ignore[override]
         @memoized_property
         def property(self): ...
 

stubs/SQLAlchemy/sqlalchemy/sql/crud.pyi

@@ -1,10 +1,13 @@
-from typing import Any
+from typing import Any, Generic, TypeVar
 
 from . import elements
+from .operators import ColumnOperators
+
+_T = TypeVar("_T")
 
 REQUIRED: Any
 
-class _multiparam_column(elements.ColumnElement[Any]):
+class _multiparam_column(elements.ColumnElement[_T], Generic[_T]):
     index: Any
     key: Any
     original: Any
@@ -12,4 +15,4 @@ class _multiparam_column(elements.ColumnElement[Any]):
     type: Any
     def __init__(self, original, index) -> None: ...
     def compare(self, other, **kw) -> None: ...
-    def __eq__(self, other): ...
+    def __eq__(self, other) -> ColumnOperators[_T]: ...  # type: ignore[override]

stubs/SQLAlchemy/sqlalchemy/sql/lambdas.pyi

@@ -1,9 +1,11 @@
-from typing import Any
+from typing import Any, Generic, TypeVar
 
 from . import elements, roles
 from .base import Options
 from .operators import ColumnOperators
 
+_T = TypeVar("_T")
+
 class LambdaOptions(Options):
     enable_tracking: bool
     track_closure_variables: bool
@@ -94,7 +96,7 @@ class AnalyzedFunction:
     closure_bindparams: Any
     def __init__(self, analyzed_code, lambda_element, apply_propagate_attrs, fn) -> None: ...
 
-class PyWrapper(ColumnOperators[Any]):
+class PyWrapper(ColumnOperators[_T], Generic[_T]):
     fn: Any
     track_bound_values: Any
     def __init__(
@@ -108,6 +110,6 @@ class PyWrapper(ColumnOperators[Any]):
     def __nonzero__(self): ...
     def __getattribute__(self, key): ...
     def __iter__(self): ...
-    def __getitem__(self, key): ...
+    def __getitem__(self, key) -> ColumnOperators[_T]: ...
 
 def insp(lmb): ...

stubs/SQLAlchemy/sqlalchemy/sql/operators.pyi

@@ -12,8 +12,8 @@ class Operators:
     def __invert__(self): ...
     def op(self, opstring, precedence: int = ..., is_comparison: bool = ..., return_type: Any | None = ...): ...
     def bool_op(self, opstring, precedence: int = ...): ...
-    def operate(self, op, *other, **kwargs) -> None: ...
-    def reverse_operate(self, op, other, **kwargs) -> None: ...
+    def operate(self, op, *other, **kwargs): ...
+    def reverse_operate(self, op, other, **kwargs): ...
 
 class custom_op:
     __name__: str
@@ -33,68 +33,68 @@ class custom_op:
         eager_grouping: bool = ...,
     ) -> None: ...
     def __eq__(self, other): ...
-    def __hash__(self): ...
+    def __hash__(self) -> int: ...
     def __call__(self, left, right, **kw): ...
 
 class ColumnOperators(Operators, Generic[_T]):
     timetuple: Any
-    def __lt__(self, other: _T | ColumnOperators[_T] | None): ...
-    def __le__(self, other: _T | ColumnOperators[_T] | None): ...
-    __hash__: Any
-    def __eq__(self, other: _T | ColumnOperators[_T] | None): ...  # type: ignore[override]
-    def __ne__(self, other: _T | ColumnOperators[_T] | None): ...  # type: ignore[override]
-    def is_distinct_from(self, other): ...
-    def is_not_distinct_from(self, other): ...
-    isnot_distinct_from = is_not_distinct_from
-    def __gt__(self, other: _T | ColumnOperators[_T] | None): ...
-    def __ge__(self, other: _T | ColumnOperators[_T] | None): ...
-    def __neg__(self): ...
-    def __contains__(self, other): ...
-    def __getitem__(self, index: int): ...
-    def __lshift__(self, other): ...
-    def __rshift__(self, other): ...
-    def concat(self, other: _T | ColumnOperators[_T] | None): ...
-    def like(self, other: _T, escape: str | None = ...): ...
-    def ilike(self, other: _T, escape: str | None = ...): ...
-    def in_(self, other: Container[_T] | Iterable[_T]): ...
-    def not_in(self, other: Container[_T] | Iterable[_T]): ...
-    notin_ = not_in
-    def not_like(self, other: _T, escape: str | None = ...): ...
-    notlike = not_like
-    def not_ilike(self, other: _T, escape: str | None = ...): ...
-    notilike = not_ilike
-    def is_(self, other: _T): ...
-    def is_not(self, other: _T): ...
-    isnot = is_not
-    def startswith(self, other: str, **kwargs): ...
-    def endswith(self, other: str, **kwargs): ...
-    def contains(self, other: str, **kwargs): ...
-    def match(self, other: str, **kwargs): ...
-    def regexp_match(self, pattern, flags: Any | None = ...): ...
-    def regexp_replace(self, pattern, replacement, flags: Any | None = ...): ...
-    def desc(self): ...
-    def asc(self): ...
-    def nulls_first(self): ...
-    nullsfirst: Any
-    def nulls_last(self): ...
-    nullslast: Any
-    def collate(self, collation): ...
-    def __radd__(self, other): ...
-    def __rsub__(self, other): ...
-    def __rmul__(self, other): ...
-    def __rdiv__(self, other): ...
-    def __rmod__(self, other): ...
-    def between(self, cleft, cright, symmetric: bool = ...): ...
-    def distinct(self): ...
-    def any_(self): ...
-    def all_(self): ...
-    def __add__(self, other): ...
-    def __sub__(self, other): ...
-    def __mul__(self, other): ...
-    def __div__(self, other): ...
-    def __mod__(self, other): ...
-    def __truediv__(self, other): ...
-    def __rtruediv__(self, other): ...
+    def __lt__(self, other: _T | ColumnOperators[_T] | None) -> ColumnOperators[_T]: ...
+    def __le__(self, other: _T | ColumnOperators[_T] | None) -> ColumnOperators[_T]: ...
+    def __hash__(self) -> int: ...
+    def __eq__(self, other: _T | ColumnOperators[_T] | None) -> ColumnOperators[_T]: ...  # type: ignore[override]
+    def __ne__(self, other: _T | ColumnOperators[_T] | None) -> ColumnOperators[_T]: ...  # type: ignore[override]
+    def is_distinct_from(self, other) -> ColumnOperators[_T]: ...
+    def is_not_distinct_from(self, other) -> ColumnOperators[_T]: ...
+    def isnot_distinct_from(self, other) -> ColumnOperators[_T]: ...
+    def __gt__(self, other: _T | ColumnOperators[_T] | None) -> ColumnOperators[_T]: ...
+    def __ge__(self, other: _T | ColumnOperators[_T] | None) -> ColumnOperators[_T]: ...
+    def __neg__(self) -> ColumnOperators[_T]: ...
+    def __contains__(self, other) -> ColumnOperators[_T]: ...
+    def __getitem__(self, index: int) -> ColumnOperators[_T]: ...
+    def __lshift__(self, other) -> ColumnOperators[_T]: ...
+    def __rshift__(self, other) -> ColumnOperators[_T]: ...
+    def concat(self, other: _T | ColumnOperators[_T] | None) -> ColumnOperators[_T]: ...
+    def like(self, other: _T, escape: str | None = ...) -> ColumnOperators[_T]: ...
+    def ilike(self, other: _T, escape: str | None = ...) -> ColumnOperators[_T]: ...
+    def in_(self, other: Container[_T] | Iterable[_T]) -> ColumnOperators[_T]: ...
+    def not_in(self, other: Container[_T] | Iterable[_T]) -> ColumnOperators[_T]: ...
+    def notin_(self, other: Container[_T] | Iterable[_T]) -> ColumnOperators[_T]: ...
+    def not_like(self, other: _T, escape: str | None = ...) -> ColumnOperators[_T]: ...
+    def notlike(self, other: _T, escape: str | None = ...) -> ColumnOperators[_T]: ...
+    def not_ilike(self, other: _T, escape: str | None = ...) -> ColumnOperators[_T]: ...
+    def notilike(self, other: _T, escape: str | None = ...) -> ColumnOperators[_T]: ...
+    def is_(self, other: _T) -> ColumnOperators[_T]: ...
+    def is_not(self, other: _T) -> ColumnOperators[_T]: ...
+    def isnot(self, other: _T) -> ColumnOperators[_T]: ...
+    def startswith(self, other: str, **kwargs) -> ColumnOperators[_T]: ...
+    def endswith(self, other: str, **kwargs) -> ColumnOperators[_T]: ...
+    def contains(self, other: str, **kwargs) -> ColumnOperators[_T]: ...
+    def match(self, other: str, **kwargs) -> ColumnOperators[_T]: ...
+    def regexp_match(self, pattern, flags: Any | None = ...) -> ColumnOperators[_T]: ...
+    def regexp_replace(self, pattern, replacement, flags: Any | None = ...) -> ColumnOperators[_T]: ...
+    def desc(self) -> ColumnOperators[_T]: ...
+    def asc(self) -> ColumnOperators[_T]: ...
+    def nulls_first(self) -> ColumnOperators[_T]: ...
+    def nullsfirst(self) -> ColumnOperators[_T]: ...
+    def nulls_last(self) -> ColumnOperators[_T]: ...
+    def nullslast(self) -> ColumnOperators[_T]: ...
+    def collate(self, collation) -> ColumnOperators[_T]: ...
+    def __radd__(self, other) -> ColumnOperators[_T]: ...
+    def __rsub__(self, other) -> ColumnOperators[_T]: ...
+    def __rmul__(self, other) -> ColumnOperators[_T]: ...
+    def __rdiv__(self, other) -> ColumnOperators[_T]: ...
+    def __rmod__(self, other) -> ColumnOperators[_T]: ...
+    def between(self, cleft, cright, symmetric: bool = ...) -> ColumnOperators[_T]: ...
+    def distinct(self) -> ColumnOperators[_T]: ...
+    def any_(self) -> ColumnOperators[_T]: ...
+    def all_(self) -> ColumnOperators[_T]: ...
+    def __add__(self, other) -> ColumnOperators[_T]: ...
+    def __sub__(self, other) -> ColumnOperators[_T]: ...
+    def __mul__(self, other) -> ColumnOperators[_T]: ...
+    def __div__(self, other) -> ColumnOperators[_T]: ...
+    def __mod__(self, other) -> ColumnOperators[_T]: ...
+    def __truediv__(self, other) -> ColumnOperators[_T]: ...
+    def __rtruediv__(self, other) -> ColumnOperators[_T]: ...
 
 def commutative_op(fn): ...
 def comparison_op(fn): ...
@@ -183,11 +183,11 @@ nullslast_op = nulls_last_op
 
 def json_getitem_op(a, b) -> None: ...
 def json_path_getitem_op(a, b) -> None: ...
-def is_comparison(op): ...
-def is_commutative(op): ...
-def is_ordering_modifier(op): ...
-def is_natural_self_precedent(op): ...
-def is_boolean(op): ...
+def is_comparison(op) -> bool: ...
+def is_commutative(op) -> bool: ...
+def is_ordering_modifier(op) -> bool: ...
+def is_natural_self_precedent(op) -> bool: ...
+def is_boolean(op) -> bool: ...
 def mirror(op): ...
-def is_associative(op): ...
-def is_precedent(operator, against): ...
+def is_associative(op) -> bool: ...
+def is_precedent(operator, against) -> bool: ...

stubs/SQLAlchemy/sqlalchemy/sql/sqltypes.pyi

@@ -1,6 +1,7 @@
-from typing import Any
+from typing import Any, Generic, TypeVar
 
 from .base import SchemaEventTarget
+from .operators import ColumnOperators
 from .traversals import HasCacheKey
 from .type_api import (
     Emulated as Emulated,
@@ -11,17 +12,19 @@ from .type_api import (
     to_instance as to_instance,
 )
 
+_T = TypeVar("_T")
+
 class _LookupExpressionAdapter:
-    class Comparator(TypeEngine.Comparator): ...
+    class Comparator(TypeEngine.Comparator[Any]): ...
     comparator_factory: Any
 
 class Concatenable:
-    class Comparator(TypeEngine.Comparator): ...
+    class Comparator(TypeEngine.Comparator[_T], Generic[_T]): ...
     comparator_factory: Any
 
 class Indexable:
-    class Comparator(TypeEngine.Comparator):
-        def __getitem__(self, index): ...
+    class Comparator(TypeEngine.Comparator[_T], Generic[_T]):
+        def __getitem__(self, index) -> ColumnOperators[_T]: ...
     comparator_factory: Any
 
 class String(Concatenable, TypeEngine):
@@ -163,7 +166,7 @@ class Enum(Emulated, String, SchemaType):
     @property
     def native(self): ...
 
-    class Comparator(Concatenable.Comparator): ...
+    class Comparator(Concatenable.Comparator[Any]): ...
     comparator_factory: Any
     def as_generic(self, allow_nulltype: bool = ...): ...
     def adapt_to_emulated(self, impltype, **kw): ...
@@ -234,7 +237,7 @@ class JSON(Indexable, TypeEngine):
     class JSONStrIndexType(JSONIndexType): ...
     class JSONPathType(JSONElementType): ...
 
-    class Comparator(Indexable.Comparator, Concatenable.Comparator):
+    class Comparator(Indexable.Comparator[Any], Concatenable.Comparator[Any]):
         def as_boolean(self): ...
         def as_string(self): ...
         def as_integer(self): ...
@@ -255,8 +258,8 @@ class ARRAY(SchemaEventTarget, Indexable, Concatenable, TypeEngine):
     __visit_name__: str
     zero_indexes: bool
 
-    class Comparator(Indexable.Comparator, Concatenable.Comparator):
-        def contains(self, *arg, **kw) -> None: ...
+    class Comparator(Indexable.Comparator[_T], Concatenable.Comparator[_T], Generic[_T]):
+        def contains(self, *arg, **kw) -> ColumnOperators[_T]: ...
         def any(self, other, operator: Any | None = ...): ...
         def all(self, other, operator: Any | None = ...): ...
     comparator_factory: Any
@@ -346,7 +349,7 @@ class NullType(TypeEngine):
     __visit_name__: str
     def literal_processor(self, dialect): ...
 
-    class Comparator(TypeEngine.Comparator): ...
+    class Comparator(TypeEngine.Comparator[Any]): ...
     comparator_factory: Any
 
 class TableValueType(HasCacheKey, TypeEngine):

stubs/SQLAlchemy/sqlalchemy/sql/type_api.pyi

@@ -1,10 +1,12 @@
-from typing import Any
+from typing import Any, Generic, TypeVar
 
 from .. import util
 from . import operators
 from .base import SchemaEventTarget
 from .visitors import Traversible, TraversibleType
 
+_T = TypeVar("_T")
+
 BOOLEANTYPE: Any
 INTEGERTYPE: Any
 NULLTYPE: Any
@@ -14,7 +16,7 @@ INDEXABLE: Any
 TABLEVALUE: Any
 
 class TypeEngine(Traversible):
-    class Comparator(operators.ColumnOperators):
+    class Comparator(operators.ColumnOperators[_T], Generic[_T]):
         default_comparator: Any
         def __clause_element__(self): ...
         expr: Any