Fix all fixable stubtest_allowlist entries in SQLAlchemy (#9596)

Co-authored-by: Alex Waygood <[email protected]>
Co-authored-by: Jelle Zijlstra <[email protected]>

Author: 3 people

stubs/SQLAlchemy/@tests/stubtest_allowlist.txt

@@ -0,0 +1,60 @@
+from __future__ import annotations
+
+from typing_extensions import assert_type
+
+from sqlalchemy.orm.strategy_options import (
+    Load,
+    contains_eager,
+    defaultload,
+    defer,
+    immediateload,
+    joinedload,
+    lazyload,
+    load_only,
+    loader_option,
+    noload,
+    raiseload,
+    selectin_polymorphic,
+    selectinload,
+    subqueryload,
+    undefer,
+    undefer_group,
+    with_expression,
+)
+
+
+def fn(loadopt: Load, *args: object) -> loader_option:
+    return loader_option()
+
+
+# Testing that the function and return type of function are actually all instances of "loader_option"
+assert_type(contains_eager, loader_option)
+assert_type(contains_eager(fn), loader_option)
+assert_type(load_only, loader_option)
+assert_type(load_only(fn), loader_option)
+assert_type(joinedload, loader_option)
+assert_type(joinedload(fn), loader_option)
+assert_type(subqueryload, loader_option)
+assert_type(subqueryload(fn), loader_option)
+assert_type(selectinload, loader_option)
+assert_type(selectinload(fn), loader_option)
+assert_type(lazyload, loader_option)
+assert_type(lazyload(fn), loader_option)
+assert_type(immediateload, loader_option)
+assert_type(immediateload(fn), loader_option)
+assert_type(noload, loader_option)
+assert_type(noload(fn), loader_option)
+assert_type(raiseload, loader_option)
+assert_type(raiseload(fn), loader_option)
+assert_type(defaultload, loader_option)
+assert_type(defaultload(fn), loader_option)
+assert_type(defer, loader_option)
+assert_type(defer(fn), loader_option)
+assert_type(undefer, loader_option)
+assert_type(undefer(fn), loader_option)
+assert_type(undefer_group, loader_option)
+assert_type(undefer_group(fn), loader_option)
+assert_type(with_expression, loader_option)
+assert_type(with_expression(fn), loader_option)
+assert_type(selectin_polymorphic, loader_option)
+assert_type(selectin_polymorphic(fn), loader_option)

stubs/SQLAlchemy/@tests/test_cases/check_loader_option.py

@@ -0,0 +1,67 @@
+from __future__ import annotations
+
+from _typeshed.dbapi import DBAPIConnection
+from typing import cast
+
+from sqlalchemy.engine.base import Engine
+from sqlalchemy.engine.default import DefaultDialect
+from sqlalchemy.engine.url import URL
+from sqlalchemy.pool.base import Pool
+from sqlalchemy.testing import config as ConfigModule
+from sqlalchemy.testing.provision import (
+    configure_follower,
+    create_db,
+    drop_all_schema_objects_post_tables,
+    drop_all_schema_objects_pre_tables,
+    drop_db,
+    follower_url_from_main,
+    generate_driver_url,
+    get_temp_table_name,
+    post_configure_engine,
+    prepare_for_drop_tables,
+    register,
+    run_reap_dbs,
+    set_default_schema_on_connection,
+    stop_test_class_outside_fixtures,
+    temp_table_keyword_args,
+    update_db_opts,
+)
+from sqlalchemy.util import immutabledict
+
+url = URL("", "", "", "", 0, "", immutabledict())
+engine = Engine(Pool(lambda: cast(DBAPIConnection, object())), DefaultDialect(), "")
+config = cast(ConfigModule.Config, object())
+unused = None
+
+
+class Foo:
+    pass
+
+
+# Test that the decorator changes the first parameter to "cfg: str | URL | _ConfigProtocol"
+@register.init
+def no_args(__foo: Foo) -> None:
+    pass
+
+
+no_args(cfg="")
+no_args(cfg=url)
+no_args(cfg=config)
+
+# Test pre-decorated functions
+generate_driver_url(url, "", "")
+drop_all_schema_objects_pre_tables(url, unused)
+drop_all_schema_objects_post_tables(url, unused)
+create_db(url, engine, unused)
+drop_db(url, engine, unused)
+update_db_opts(url, unused)
+post_configure_engine(url, unused, unused)
+follower_url_from_main(url, "")
+configure_follower(url, unused)
+run_reap_dbs(url, unused)
+temp_table_keyword_args(url, engine)
+prepare_for_drop_tables(url, unused)
+stop_test_class_outside_fixtures(url, unused, type)
+get_temp_table_name(url, unused, "")
+set_default_schema_on_connection(ConfigModule, unused, unused)
+set_default_schema_on_connection(config, unused, unused)

stubs/SQLAlchemy/@tests/test_cases/check_register.py

@@ -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

@@ -22,6 +22,7 @@ class _URLTuple(NamedTuple):
 _Query: TypeAlias = Mapping[str, str | Sequence[str]] | Sequence[tuple[str, str | Sequence[str]]]
 
 class URL(_URLTuple):
+    def __new__(self, *arg, **kw) -> Self | URL: ...
     @classmethod
     def create(
         cls,

stubs/SQLAlchemy/sqlalchemy/engine/url.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 _typeshed import Incomplete, SupportsKeysAndGetItem
+from collections.abc import Iterable
+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,41 @@ class CollectionAdapter:
     def fire_remove_event(self, item, initiator: Incomplete | None = None) -> None: ...
     def fire_pre_remove_event(self, initiator: Incomplete | None = 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) -> None: ...
-    def remove(self, value, _sa_initiator: Incomplete | None = None) -> None: ...
+    def set(self, value: _VT, _sa_initiator: Event | Literal[False] | None = None) -> None: ...
+    def remove(self, value: _VT, _sa_initiator: Event | Literal[False] | None = 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: ...
+    @overload
+    def update(self, __other: SupportsKeysAndGetItem[_KT, _VT] = ..., **kwargs: _VT) -> None: ...
+    @overload
+    def update(self, __other: Iterable[tuple[_KT, _VT]] = ..., **kwargs: _VT) -> None: ...

stubs/SQLAlchemy/sqlalchemy/orm/collections.pyi

@@ -13,7 +13,7 @@ _DeclT = TypeVar("_DeclT", bound=type[_DeclarativeBase])
 
 # Dynamic class as created by registry.generate_base() via DeclarativeMeta
 # or another metaclass. This class does not exist at runtime.
-class _DeclarativeBase(Any):  # super classes are dynamic
+class _DeclarativeBase(Any):  # type: ignore[misc]  # super classes are dynamic
     registry: ClassVar[registry]
     metadata: ClassVar[MetaData]
     __abstract__: ClassVar[bool]