Fix all fixable stubtest_allowlist entries in SQLAlchemy (#9596)

Co-authored-by: Alex Waygood <Alex.Waygood@Gmail.com>
Co-authored-by: Jelle Zijlstra <jelle.zijlstra@gmail.com>

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

@@ -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_register.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/sqlalchemy/dialects/mssql/base.pyi

@@ -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/postgresql/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/engine/interfaces.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/url.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/event/base.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/orm/collections.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 MappedCollection(dict[Any, Any]):
+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[_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/decl_api.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]

stubs/SQLAlchemy/sqlalchemy/orm/strategy_options.pyi

@@ -1,5 +1,7 @@
 from _typeshed import Incomplete
+from collections.abc import Callable
 from typing import Any
+from typing_extensions import Self
 
 from ..sql.base import Generative
 from .interfaces import LoaderOption
@@ -17,27 +19,27 @@ class Load(Generative, LoaderOption):
     propagate_to_loaders: bool
     def process_compile_state_replaced_entities(self, compile_state, mapper_entities) -> None: ...
     def process_compile_state(self, compile_state) -> None: ...
-    def options(self, *opts) -> None: ...
-    def set_relationship_strategy(self, attr, strategy, propagate_to_loaders: bool = True) -> None: ...
-    def set_column_strategy(self, attrs, strategy, opts: Incomplete | None = None, opts_only: bool = False) -> None: ...
-    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 = None): ...
-    def load_only(self, *attrs): ...
-    def joinedload(self, attr, innerjoin: Incomplete | None = 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 = False): ...
-    def defaultload(self, attr): ...
-    def defer(self, key, raiseload: bool = False): ...
-    def undefer(self, key): ...
-    def undefer_group(self, name): ...
-    def with_expression(self, key, expression): ...
-    def selectin_polymorphic(self, classes): ...
+    def options(self, *opts) -> Self: ...
+    def set_relationship_strategy(self, attr, strategy, propagate_to_loaders: bool = True) -> Self: ...
+    def set_column_strategy(self, attrs, strategy, opts: Incomplete | None = None, opts_only: bool = False) -> Self: ...
+    def set_generic_strategy(self, attrs, strategy) -> Self: ...
+    def set_class_strategy(self, strategy, opts) -> Self: ...
+    # Added dynamically at runtime
+    def contains_eager(loadopt, attr, alias: Incomplete | None = None) -> Self: ...
+    def load_only(loadopt, *attrs) -> Self: ...
+    def joinedload(loadopt, attr, innerjoin: Incomplete | None = None) -> Self: ...
+    def subqueryload(loadopt, attr) -> Self: ...
+    def selectinload(loadopt, attr) -> Self: ...
+    def lazyload(loadopt, attr) -> Self: ...
+    def immediateload(loadopt, attr) -> Self: ...
+    def noload(loadopt, attr) -> Self: ...
+    def raiseload(loadopt, attr, sql_only: bool = False) -> Self: ...
+    def defaultload(loadopt, attr) -> Self: ...
+    def defer(loadopt, key, raiseload: bool = False) -> Self: ...
+    def undefer(loadopt, key) -> Self: ...
+    def undefer_group(loadopt, name) -> Self: ...
+    def with_expression(loadopt, key, expression) -> Self: ...
+    def selectin_polymorphic(loadopt, classes) -> Self: ...
 
 class _UnboundLoad(Load):
     path: Any
@@ -45,22 +47,39 @@ class _UnboundLoad(Load):
     def __init__(self) -> None: ...
 
 class loader_option:
-    name: Any
-    fn: Any
-    def __call__(self, fn): ...
+    name: str
+    # The first parameter of this Callable should always be `loadopt: Load`
+    fn: Callable[..., loader_option]
+    def __call__(self, fn: Callable[..., loader_option]) -> Self: ...
 
-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): ...
+# loader_option instances that can be used to dynamically add methods to Load at runtime
+@loader_option()
+def contains_eager(loadopt: Load, attr, alias: Incomplete | None = ...) -> loader_option: ...
+@loader_option()
+def load_only(loadopt: Load, *attrs) -> loader_option: ...
+@loader_option()
+def joinedload(loadopt, attr, innerjoin=None): ...
+@loader_option()
+def subqueryload(loadopt: Load, attr) -> loader_option: ...
+@loader_option()
+def selectinload(loadopt: Load, attr) -> loader_option: ...
+@loader_option()
+def lazyload(loadopt: Load, attr) -> loader_option: ...
+@loader_option()
+def immediateload(loadopt: Load, attr) -> loader_option: ...
+@loader_option()
+def noload(loadopt: Load, attr) -> loader_option: ...
+@loader_option()
+def raiseload(loadopt: Load, attr, sql_only: bool = ...) -> loader_option: ...
+@loader_option()
+def defaultload(loadopt: Load, attr) -> loader_option: ...
+@loader_option()
+def defer(loadopt: Load, key, raiseload: bool = ...) -> loader_option: ...
+@loader_option()
+def undefer(loadopt: Load, key) -> loader_option: ...
+@loader_option()
+def undefer_group(loadopt: Load, name) -> loader_option: ...
+@loader_option()
+def with_expression(loadopt: Load, key) -> loader_option: ...
+@loader_option()
+def selectin_polymorphic(loadopt: Load, classes) -> loader_option: ...

stubs/SQLAlchemy/sqlalchemy/pool/base.pyi

@@ -1,4 +1,6 @@
 from _typeshed import Incomplete
+from _typeshed.dbapi import DBAPIConnection
+from collections.abc import Callable
 from typing import Any
 
 from .. import log
@@ -24,7 +26,7 @@ class Pool(log.Identified):
     echo: Any
     def __init__(
         self,
-        creator,
+        creator: Callable[[], DBAPIConnection],
         recycle: int = -1,
         echo: Incomplete | None = None,
         logging_name: Incomplete | None = None,

stubs/SQLAlchemy/sqlalchemy/sql/base.pyi

@@ -73,7 +73,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/compiler.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/functions.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

stubs/SQLAlchemy/sqlalchemy/sql/type_api.pyi

@@ -55,7 +55,7 @@ class VisitableCheckKWArg(util.EnsureKWArgType, TraversibleType): ...
 class ExternalType:
     cache_ok: Any
 
-class UserDefinedType:
+class UserDefinedType(ExternalType, TypeEngine, metaclass=VisitableCheckKWArg):
     __visit_name__: str
     ensure_kwarg: str
     def coerce_compared_value(self, op, value): ...

stubs/SQLAlchemy/sqlalchemy/sql/visitors.pyi

@@ -3,14 +3,14 @@ from typing import Any
 class TraversibleType(type):
     def __init__(cls, clsname, bases, clsdict) -> None: ...
 
-class Traversible:
+class Traversible(metaclass=TraversibleType):
     def __class_getitem__(cls, key): ...
     def get_children(self, omit_attrs=(), **kw): ...
 
 class _InternalTraversalType(type):
     def __init__(cls, clsname, bases, clsdict) -> None: ...
 
-class InternalTraversal:
+class InternalTraversal(metaclass=_InternalTraversalType):
     def dispatch(self, visit_symbol): ...
     def run_generated_dispatch(self, target, internal_dispatch, generate_dispatcher_name): ...
     def generate_dispatch(self, target_cls, internal_dispatch, generate_dispatcher_name): ...

stubs/SQLAlchemy/sqlalchemy/testing/config.pyi

@@ -1,12 +1,13 @@
-from typing import Any
+from _typeshed import Incomplete
+from typing import Any, Protocol
 
-requirements: Any
-db: Any
-db_url: Any
-db_opts: Any
-file_config: Any
-test_schema: Any
-test_schema_2: Any
+requirements: Incomplete | None
+db: Incomplete | None
+db_url: Incomplete | None
+db_opts: Incomplete | None
+file_config: Incomplete | None
+test_schema: str | None
+test_schema_2: str | None
 any_async: bool
 ident: str
 
@@ -16,17 +17,26 @@ def fixture(*arg, **kw): ...
 def get_current_test_name(): ...
 def mark_base_test_class(): ...
 
+# Matches the intersection of the config module and the Config class
+class _ConfigProtocol(Protocol):  # noqa: Y046
+    db: Incomplete
+    db_opts: Incomplete
+    file_config: Incomplete
+    test_schema: Any  # AnyOf[str, None]
+    test_schema_2: Any  # AnyOf[str, None]
+    def skip_test(self, msg) -> None: ...
+
 class Config:
-    db: Any
-    db_opts: Any
-    options: Any
-    file_config: Any
+    db: Incomplete
+    db_opts: Incomplete
+    options: Incomplete
+    file_config: Incomplete
     test_schema: str
     test_schema_2: str
-    is_async: Any
+    is_async: Incomplete
     def __init__(self, db, db_opts, options, file_config) -> None: ...
     @classmethod
-    def register(cls, db, db_opts, options, file_config): ...
+    def register(cls, db, db_opts, options, file_config) -> Config: ...
     @classmethod
     def set_as_current(cls, config, namespace) -> None: ...
     @classmethod

stubs/SQLAlchemy/sqlalchemy/testing/provision.pyi

@@ -1,33 +1,63 @@
-from typing import Any
+from _typeshed import Incomplete
+from collections.abc import Callable
+from logging import Logger
+from typing import Any, Generic, NoReturn, TypeVar
+from typing_extensions import Self, TypeAlias
 
-log: Any
-FOLLOWER_IDENT: Any
+from ..engine.interfaces import Connectable
+from ..engine.url import URL
+from .config import Config, _ConfigProtocol
 
-class register:
-    fns: Any
-    @classmethod
-    def init(cls, fn): ...
-    def for_db(self, *dbnames): ...
-    def __call__(self, cfg, *arg): ...
+_Unused: TypeAlias = object
+_S = TypeVar("_S", bound=str)
+_U = TypeVar("_U", bound=URL)
+_F = TypeVar("_F", bound=Callable[..., str | URL | None])
+
+log: Logger
+FOLLOWER_IDENT: Incomplete | None
 
 def create_follower_db(follower_ident) -> None: ...
-def setup_config(db_url, options, file_config, follower_ident): ...
+def setup_config(db_url, options, file_config, follower_ident) -> Config: ...
 def drop_follower_db(follower_ident) -> None: ...
 def generate_db_urls(db_urls, extra_drivers) -> None: ...
-def generate_driver_url(url, driver, query_str): ...
-def drop_all_schema_objects_pre_tables(cfg, eng) -> None: ...
-def drop_all_schema_objects_post_tables(cfg, eng) -> None: ...
 def drop_all_schema_objects(cfg, eng) -> None: ...
-def create_db(cfg, eng, ident) -> None: ...
-def drop_db(cfg, eng, ident) -> None: ...
-def update_db_opts(cfg, db_opts) -> None: ...
-def post_configure_engine(url, engine, follower_ident) -> None: ...
-def follower_url_from_main(url, ident): ...
-def configure_follower(cfg, ident) -> None: ...
-def run_reap_dbs(url, ident) -> None: ...
 def reap_dbs(idents_file) -> None: ...
-def temp_table_keyword_args(cfg, eng) -> None: ...
-def prepare_for_drop_tables(config, connection) -> None: ...
-def stop_test_class_outside_fixtures(config, db, testcls) -> None: ...
-def get_temp_table_name(cfg, eng, base_name): ...
-def set_default_schema_on_connection(cfg, dbapi_connection, schema_name) -> None: ...
+
+class register(Generic[_F]):
+    fns: dict[str, _F]
+    @classmethod
+    def init(cls, fn: _F) -> Self: ...
+    def for_db(self, *dbnames: str) -> Callable[[_F], Self]: ...
+    # Impossible to specify the args from the generic Callable in the current type system
+    def __call__(self, cfg: str | URL | _ConfigProtocol, *arg: Any) -> str | URL | None: ...  # AnyOf[str | URL | None]
+
+@register.init
+def generate_driver_url(url: _U, driver: str, query_str: str) -> _U | None: ...
+@register.init
+def drop_all_schema_objects_pre_tables(cfg: _Unused, eng: _Unused) -> None: ...
+@register.init
+def drop_all_schema_objects_post_tables(cfg: _Unused, eng: _Unused) -> None: ...
+@register.init
+def create_db(cfg: _Unused, eng: Connectable, ident: _Unused) -> NoReturn: ...
+@register.init
+def drop_db(cfg: _Unused, eng: Connectable, ident: _Unused) -> NoReturn: ...
+@register.init
+def update_db_opts(db_url: _Unused, db_opts: _Unused) -> None: ...
+@register.init
+def post_configure_engine(url: _Unused, engine: _Unused, follower_ident: _Unused) -> None: ...
+@register.init
+def follower_url_from_main(url: _U, ident: str) -> _U: ...
+@register.init
+def configure_follower(cfg: _Unused, ident: _Unused) -> None: ...
+@register.init
+def run_reap_dbs(url: _Unused, ident: _Unused) -> None: ...
+@register.init
+def temp_table_keyword_args(cfg: _Unused, eng: Connectable) -> NoReturn: ...
+@register.init
+def prepare_for_drop_tables(config: _Unused, connection: _Unused) -> None: ...
+@register.init
+def stop_test_class_outside_fixtures(config: _Unused, db: _Unused, testcls: _Unused) -> None: ...
+@register.init  # type: ignore[type-var]  # False-positive, _S is bound to str
+def get_temp_table_name(cfg: _Unused, eng: _Unused, base_name: _S) -> _S: ...
+@register.init
+def set_default_schema_on_connection(cfg: _ConfigProtocol, dbapi_connection: _Unused, schema_name: _Unused) -> NoReturn: ...

stubs/SQLAlchemy/sqlalchemy/util/__init__.pyi

@@ -59,6 +59,7 @@ from .compat import (
     has_refcount_gc as has_refcount_gc,
     inspect_getfullargspec as inspect_getfullargspec,
     int_types as int_types,
+    is64bit as is64bit,
     iterbytes as iterbytes,
     itertools_filter as itertools_filter,
     itertools_filterfalse as itertools_filterfalse,
@@ -76,6 +77,9 @@ from .compat import (
     py37 as py37,
     py38 as py38,
     py39 as py39,
+    py310 as py310,
+    py311 as py311,
+    py312 as py312,
     pypy as pypy,
     quote_plus as quote_plus,
     raise_ as raise_,

stubs/SQLAlchemy/sqlalchemy/util/_collections.pyi

@@ -140,7 +140,7 @@ class IdentitySet:
     def __hash__(self) -> int: ...
 
 class WeakSequence:
-    def __init__(self, __elements=...) -> None: ...
+    def __init__(self, _WeakSequence__elements: Iterable[Incomplete] = ()) -> None: ...
     def append(self, item) -> None: ...
     def __len__(self) -> int: ...
     def __iter__(self): ...

stubs/SQLAlchemy/sqlalchemy/util/compat.pyi

@@ -1,19 +1,20 @@
-import builtins
-import collections
-import contextlib
 import itertools
 import operator
 import pickle as pickle
+import sys
 import threading as threading
 from _typeshed import Incomplete, Unused
 from abc import ABC as ABC
+from builtins import callable as callable, next as next
+from collections import namedtuple as namedtuple  # noqa: Y024  # Actual import
+from contextlib import contextmanager as contextmanager
 from datetime import timezone as timezone
 from functools import reduce as reduce
-from io import BytesIO as BytesIO, StringIO as StringIO
+from io import BytesIO, StringIO as StringIO
 from itertools import zip_longest as zip_longest
 from time import perf_counter as perf_counter
 from typing import TYPE_CHECKING as TYPE_CHECKING, Any, NamedTuple
-from typing_extensions import Literal
+from typing_extensions import Final
 from urllib.parse import (
     parse_qsl as parse_qsl,
     quote as quote,
@@ -24,21 +25,28 @@ from urllib.parse import (
 
 byte_buffer = BytesIO
 
+py312: bool
+py311: bool
+py310: bool
 py39: bool
 py38: bool
 py37: bool
-py3k: Literal[True]
-py2k: Literal[False]
+py3k: Final = True
+py2k: Final = False
 pypy: bool
 cpython: bool
-win32: bool
-osx: bool
+if sys.platform == "win32":
+    win32: Final = True
+else:
+    win32: Final = False
+if sys.platform == "darwin":
+    osx: Final = True
+else:
+    osx: Final = False
 arm: bool
+is64bit: bool
 has_refcount_gc: bool
-contextmanager = contextlib.contextmanager
 dottedgetter = operator.attrgetter
-namedtuple = collections.namedtuple  # noqa: Y024
-next = builtins.next
 
 class FullArgSpec(NamedTuple):
     args: Any
@@ -82,10 +90,6 @@ def raise_(
 ) -> None: ...
 def u(s): ...
 def ue(s): ...
-
-callable = builtins.callable
-
-def safe_bytestring(text): ...
 def inspect_formatargspec(
     args,
     varargs: Incomplete | None = None,

stubs/SQLAlchemy/sqlalchemy/util/langhelpers.pyi

@@ -34,7 +34,7 @@ class PluginLoader:
     def __init__(self, group, auto_fn: Incomplete | None = None) -> None: ...
     def clear(self) -> None: ...
     def load(self, name): ...
-    def register(self, name, modulepath, objname): ...
+    def register(self, name, modulepath, objname) -> None: ...
 
 def get_cls_kwargs(cls, _set: Incomplete | None = None): ...
 def get_func_kwargs(func): ...
@@ -133,12 +133,16 @@ class hybridmethod:
     def classlevel(self, func): ...
 
 class _symbol(int):
-    def __new__(cls, name, doc: Incomplete | None = None, canonical: Incomplete | None = None): ...
+    def __new__(  # noqa: Y034  # Explicitly instantiates _symbol
+        self, name, doc: Incomplete | None = None, canonical: Incomplete | None = None
+    ) -> _symbol: ...
     def __reduce__(self): ...
 
 class symbol:
     symbols: Any
-    def __new__(cls, name, doc: Incomplete | None = None, canonical: Incomplete | None = None): ...
+    def __new__(  # type: ignore[misc]  # Explicitly instantiates _symbol
+        cls, name, doc: Incomplete | None = None, canonical: Incomplete | None = None
+    ) -> _symbol: ...
     @classmethod
     def parse_user_argument(cls, arg, choices, name, resolve_symbol_names: bool = False): ...