Fix most stubtest_allowlist entries in SQLAlchemy

Author: Avasam

stubs/SQLAlchemy/@tests/stubtest_allowlist.txt

@@ -197,7 +197,6 @@ class MSExecutionContext(default.DefaultExecutionContext):
     @property
     def rowcount(self): ...
     def handle_dbapi_exception(self, e) -> None: ...
-    def get_result_cursor_strategy(self, result): ...
     def fire_sequence(self, seq, type_): ...
     def get_insert_default(self, column): ...
 

stubs/SQLAlchemy/sqlalchemy/dialects/mssql/base.pyi

@@ -163,6 +163,7 @@ class PGCompiler(compiler.SQLCompiler):
 class PGDDLCompiler(compiler.DDLCompiler):
     def get_column_specification(self, column, **kwargs): ...
     def visit_check_constraint(self, constraint): ...
+    def visit_foreign_key_constraint(self, constraint) -> str: ...  # type: ignore[override]  # Different params
     def visit_drop_table_comment(self, drop): ...
     def visit_create_enum_type(self, create): ...
     def visit_drop_enum_type(self, drop): ...

stubs/SQLAlchemy/sqlalchemy/dialects/postgresql/base.pyi

@@ -116,7 +116,6 @@ class ExecutionContext:
     def pre_exec(self) -> None: ...
     def get_out_parameter_values(self, out_param_names) -> None: ...
     def post_exec(self) -> None: ...
-    def get_result_cursor_strategy(self, result) -> None: ...
     def handle_dbapi_exception(self, e) -> None: ...
     def should_autocommit_text(self, statement) -> None: ...
     def lastrow_has_defaults(self) -> None: ...

stubs/SQLAlchemy/sqlalchemy/engine/interfaces.pyi

@@ -12,7 +12,7 @@ class _Dispatch:
 class _EventMeta(type):
     def __init__(cls, classname, bases, dict_) -> None: ...
 
-class Events:
+class Events(metaclass=_EventMeta):
     dispatch: Any
 
 class _JoinedDispatcher:

stubs/SQLAlchemy/sqlalchemy/event/base.pyi

@@ -1,5 +1,14 @@
 from _typeshed import Incomplete
-from typing import Any
+from collections.abc import Iterable, Mapping
+from typing import Any, TypeVar, overload
+from typing_extensions import Literal, SupportsIndex
+
+from ..orm.attributes import Event
+from ..util.langhelpers import _symbol, symbol
+
+_T = TypeVar("_T")
+_KT = TypeVar("_KT")
+_VT = TypeVar("_VT")
 
 class _PlainColumnGetter:
     cols: Any
@@ -81,12 +90,38 @@ class CollectionAdapter:
     def fire_remove_event(self, item, initiator: Incomplete | None = ...) -> None: ...
     def fire_pre_remove_event(self, initiator: Incomplete | None = ...) -> None: ...
 
-class InstrumentedList(list[Any]): ...
-class InstrumentedSet(set[Any]): ...
-class InstrumentedDict(dict[Any, Any]): ...
+class InstrumentedList(list[_T]):
+    def append(self, item, _sa_initiator: Event | Literal[False] | None = None) -> None: ...
+    def clear(self, index: SupportsIndex = -1) -> None: ...
+    def extend(self, iterable: Iterable[_T]) -> None: ...
+    def insert(self, index: SupportsIndex, value: _T) -> None: ...
+    def pop(self, index: SupportsIndex = -1) -> _T: ...
+    def remove(self, value: _T, _sa_initiator: Event | Literal[False] | None = None) -> None: ...
+
+class InstrumentedSet(set[_T]):
+    def add(self, value: _T, _sa_initiator: Event | Literal[False] | None = None) -> None: ...
+    def difference_update(self, value: Iterable[_T]) -> None: ...  # type: ignore[override]
+    def discard(self, value: _T, _sa_initiator: Event | Literal[False] | None = None) -> None: ...
+    def intersection_update(self, other: Iterable[_T]) -> None: ...  # type: ignore[override]
+    def remove(self, value: _T, _sa_initiator: Event | Literal[False] | None = None) -> None: ...
+    def symmetric_difference_update(self, other: Iterable[_T]) -> None: ...
+    def update(self, value: Iterable[_T]) -> None: ...  # type: ignore[override]
+
+class InstrumentedDict(dict[_KT, _VT]): ...
 
-class MappedCollection(dict[Any, Any]):
+class MappedCollection(dict[_KT, _VT]):
     keyfunc: Any
     def __init__(self, keyfunc) -> None: ...
-    def set(self, value, _sa_initiator: Incomplete | None = ...) -> None: ...
-    def remove(self, value, _sa_initiator: Incomplete | None = ...) -> None: ...
+    def set(self, value: _VT, _sa_initiator: Event | Literal[False] | None = ...) -> None: ...
+    def remove(self, value: _VT, _sa_initiator: Event | Literal[False] | None = ...) -> None: ...
+    def __delitem__(self, key: _KT, _sa_initiatorEvent: Event | Literal[False] | None = None) -> None: ...
+    def __setitem__(self, key: _KT, value: _VT, _sa_initiator: Event | Literal[False] | None = None) -> None: ...
+    @overload
+    def pop(self, key: _KT) -> _VT: ...
+    @overload
+    def pop(self, key: _KT, default: _VT | _T | _symbol | symbol = ...) -> _VT | _T: ...
+    @overload  # type: ignore[override]
+    def setdefault(self, key: _KT, default: _T) -> _VT | _T: ...
+    @overload
+    def setdefault(self, key: _KT, default: None = None) -> _VT | None: ...
+    def update(self, __other: Mapping[_KT, _VT] = ..., **kw: _VT) -> None: ...  # type: ignore[override]

stubs/SQLAlchemy/sqlalchemy/orm/collections.pyi

@@ -1,5 +1,6 @@
 from _typeshed import Incomplete
-from typing import Any
+from collections import Callable
+from typing import Any, Generic, Protocol, TypeVar
 
 from ..sql.base import Generative
 from .interfaces import LoaderOption
@@ -23,44 +24,116 @@ class Load(Generative, LoaderOption):
     def set_generic_strategy(self, attrs, strategy) -> None: ...
     def set_class_strategy(self, strategy, opts) -> None: ...
     # added dynamically at runtime
-    def contains_eager(self, attr, alias: Incomplete | None = ...): ...
-    def load_only(self, *attrs): ...
-    def joinedload(self, attr, innerjoin: Incomplete | None = ...): ...
-    def subqueryload(self, attr): ...
-    def selectinload(self, attr): ...
-    def lazyload(self, attr): ...
-    def immediateload(self, attr): ...
-    def noload(self, attr): ...
-    def raiseload(self, attr, sql_only: bool = ...): ...
-    def defaultload(self, attr): ...
-    def defer(self, key, raiseload: bool = ...): ...
-    def undefer(self, key): ...
-    def undefer_group(self, name): ...
-    def with_expression(self, key, expression): ...
-    def selectin_polymorphic(self, classes): ...
+    def contains_eager(loadopt, attr, alias: Incomplete | None = ...): ...
+    def load_only(loadopt, *attrs): ...
+    def joinedload(loadopt, attr, innerjoin: Incomplete | None = ...): ...
+    def subqueryload(loadopt, attr): ...
+    def selectinload(loadopt, attr): ...
+    def lazyload(loadopt, attr): ...
+    def immediateload(loadopt, attr): ...
+    def noload(loadopt, attr): ...
+    def raiseload(loadopt, attr, sql_only: bool = ...): ...
+    def defaultload(loadopt, attr): ...
+    def defer(loadopt, key, raiseload: bool = ...): ...
+    def undefer(loadopt, key): ...
+    def undefer_group(loadopt, name): ...
+    def with_expression(loadopt, key, expression): ...
+    def selectin_polymorphic(loadopt, classes): ...
 
 class _UnboundLoad(Load):
     path: Any
     local_opts: Any
     def __init__(self) -> None: ...
 
-class loader_option:
-    name: Any
-    fn: Any
-    def __call__(self, fn): ...
-
-def contains_eager(loadopt, attr, alias: Incomplete | None = ...): ...
-def load_only(loadopt, *attrs): ...
-def joinedload(loadopt, attr, innerjoin: Incomplete | None = ...): ...
-def subqueryload(loadopt, attr): ...
-def selectinload(loadopt, attr): ...
-def lazyload(loadopt, attr): ...
-def immediateload(loadopt, attr): ...
-def noload(loadopt, attr): ...
-def raiseload(loadopt, attr, sql_only: bool = ...): ...
-def defaultload(loadopt, attr): ...
-def defer(loadopt, key, raiseload: bool = ...): ...
-def undefer(loadopt, key): ...
-def undefer_group(loadopt, name): ...
-def with_expression(loadopt, key, expression): ...
-def selectin_polymorphic(loadopt, classes): ...
+###
+# The methods below are decorated with the class loader_option
+# They dynamically become instances of loader_option,
+# wich is callable with their original parameters.
+#
+# While both mypy and pyright's validation work, Pylance is unable to
+# show the parameters and return types.
+#
+# There is a workaround (define the method for Pylance, then reassign
+# an instance of loader_option to it for mypy, and add pyright+Flake8
+# suppressions), but it is too hacky and relies on some unsupported quirks.
+#
+# Asking Pylance to add support for these generic callables might be preferable.
+###
+
+_F = TypeVar("_F", bound=Callable[..., loader_option[Any]])
+
+class loader_option(Generic[_F]):
+    name: str
+    _dynamic: _F
+    fn: _F
+    __call__: _F  # Cheesy "__call__" definition to use the dynamic methods instead
+
+class _contains_eager(Protocol):
+    def __call__(self, loadopt, attr, alias: Incomplete | None = ...) -> loader_option[_contains_eager]: ...
+
+contains_eager: loader_option[_contains_eager]
+
+class _load_only(Protocol):
+    def __call__(self, loadopt, *attrs) -> loader_option[_load_only]: ...
+
+load_only: loader_option[_load_only]
+
+class _subqueryload(Protocol):
+    def __call__(self, loadopt, attr) -> loader_option[_subqueryload]: ...
+
+subqueryload: loader_option[_subqueryload]
+
+class _selectinload(Protocol):
+    def __call__(self, loadopt, attr) -> loader_option[_selectinload]: ...
+
+selectinload: loader_option[_selectinload]
+
+class _lazyload(Protocol):
+    def __call__(self, loadopt, attr) -> loader_option[_lazyload]: ...
+
+lazyload: loader_option[_lazyload]
+
+class _immediateload(Protocol):
+    def __call__(self, loadopt, attr) -> loader_option[_immediateload]: ...
+
+immediateload: loader_option[_immediateload]
+
+class _noload(Protocol):
+    def __call__(self, loadopt, attr) -> loader_option[_noload]: ...
+
+noload: loader_option[_noload]
+
+class _raiseload(Protocol):
+    def __call__(self, loadopt, attr, sql_only: bool = ...) -> loader_option[_raiseload]: ...
+
+raiseload: loader_option[_raiseload]
+
+class _defaultload(Protocol):
+    def __call__(self, loadopt, attr) -> loader_option[_defaultload]: ...
+
+defaultload: loader_option[_defaultload]
+
+class _defer(Protocol):
+    def __call__(self, loadopt, attr, sql_only: bool = ...) -> loader_option[_defer]: ...
+
+defer: loader_option[_defer]
+
+class _undefer(Protocol):
+    def __call__(self, loadopt, key) -> loader_option[_undefer]: ...
+
+undefer: loader_option[_undefer]
+
+class _undefer_group(Protocol):
+    def __call__(self, loadopt, name) -> loader_option[_undefer_group]: ...
+
+undefer_group: loader_option[_undefer_group]
+
+class _with_expression(Protocol):
+    def __call__(self, loadopt, key, expression) -> loader_option[_with_expression]: ...
+
+with_expression: loader_option[_with_expression]
+
+class _selectin_polymorphic(Protocol):
+    def __call__(self, loadopt, classes) -> loader_option[_selectin_polymorphic]: ...
+
+selectin_polymorphic: loader_option[_selectin_polymorphic]

stubs/SQLAlchemy/sqlalchemy/orm/strategy_options.pyi

@@ -72,7 +72,7 @@ class _MetaOptions(type):
     def __init__(cls, classname, bases, dict_) -> None: ...
     def __add__(self, other): ...
 
-class Options:
+class Options(metaclass=_MetaOptions):
     def __init__(self, **kw) -> None: ...
     def __add__(self, other): ...
     def __eq__(self, other): ...

stubs/SQLAlchemy/sqlalchemy/sql/base.pyi

@@ -1,6 +1,8 @@
 from _typeshed import Incomplete
 from typing import NamedTuple
 
+from sqlalchemy.util.langhelpers import EnsureKWArgType
+
 from ..util import memoized_property
 from . import elements
 
@@ -67,7 +69,7 @@ class Compiled:
     @property
     def params(self): ...
 
-class TypeCompiler:
+class TypeCompiler(metaclass=EnsureKWArgType):
     ensure_kwarg: str
     dialect: Incomplete
     def __init__(self, dialect) -> None: ...

stubs/SQLAlchemy/sqlalchemy/sql/compiler.pyi

@@ -87,7 +87,7 @@ class Function(FunctionElement):
 class _GenericMeta(TraversibleType):
     def __init__(cls, clsname, bases, clsdict) -> None: ...
 
-class GenericFunction:
+class GenericFunction(Function, metaclass=_GenericMeta):
     name: Incomplete
     identifier: Incomplete
     coerce_arguments: bool