Further annotate sqlalchemy.engine and collections

stubs/SQLAlchemy/@tests/stubtest_allowlist.txt

@@ -1,3 +1,6 @@
+# stub-only module
+sqlalchemy.dbapi
+
 # wrong argument name in implementation ("self" instead of "cls")
 sqlalchemy.engine.URL.__new__
 sqlalchemy.engine.url.URL.__new__
@@ -25,6 +28,8 @@ sqlalchemy.testing.util.resolve_lambda
 sqlalchemy.util.WeakSequence.__init__
 
 # not always present
+sqlalchemy.engine.Engine.logging_name  # initialized if not None
+sqlalchemy.engine.base.Engine.logging_name  # initialized if not None
 sqlalchemy.testing.util.non_refcount_gc_collect
 
 # replaced at runtime
@@ -103,6 +108,20 @@ sqlalchemy.orm.strategy_options.Load.undefer
 sqlalchemy.orm.strategy_options.Load.undefer_group
 sqlalchemy.orm.strategy_options.Load.with_expression
 
+# abstract fields not present at runtime
+sqlalchemy.engine.Transaction.connection
+sqlalchemy.engine.Transaction.is_active
+sqlalchemy.engine.base.Transaction.connection
+sqlalchemy.engine.base.Transaction.is_active
+
+# initialized to None during class construction, but overridden during __init__()
+sqlalchemy.engine.Connection.engine
+sqlalchemy.engine.base.Connection.engine
+
+# uses @memoized_property at runtime, but we use @property for compatibility
+sqlalchemy.engine.URL.normalized_query
+sqlalchemy.engine.url.URL.normalized_query
+
 # unclear problems
 sqlalchemy.sql.elements.quoted_name.lower
 sqlalchemy.sql.elements.quoted_name.upper

stubs/SQLAlchemy/sqlalchemy/cimmutabledict.pyi

@@ -1,12 +1,17 @@
-class immutabledict:
-    def __len__(self) -> int: ...
-    def __getitem__(self, __item): ...
-    def __iter__(self): ...
-    def union(self, **kwargs): ...
-    def merge_with(self, *args): ...
-    def keys(self): ...
-    def __contains__(self, __item): ...
-    def items(self): ...
-    def values(self): ...
-    def get(self, __key, __default=...): ...
-    def __reduce__(self): ...
+from _typeshed import SupportsKeysAndGetItem
+from collections.abc import Iterable
+from typing import Generic, TypeVar, overload
+
+_KT = TypeVar("_KT")
+_KT2 = TypeVar("_KT2")
+_VT = TypeVar("_VT")
+_VT2 = TypeVar("_VT2")
+
+class immutabledict(dict[_KT, _VT], Generic[_KT, _VT]):
+    @overload
+    def union(self, __dict: dict[_KT2, _VT2]) -> immutabledict[_KT | _KT2, _VT | _VT2]: ...
+    @overload
+    def union(self, __dict: None = ..., **kw: SupportsKeysAndGetItem[_KT2, _VT2]) -> immutabledict[_KT | _KT2, _VT | _VT2]: ...
+    def merge_with(
+        self, *args: SupportsKeysAndGetItem[_KT | _KT2, _VT2] | Iterable[tuple[_KT2, _VT2]] | None
+    ) -> immutabledict[_KT | _KT2, _VT | _VT2]: ...

stubs/SQLAlchemy/sqlalchemy/dbapi.pyi

@@ -0,0 +1,36 @@
+# TODO: Tempory copy of _typeshed.dbapi, until that file is available in all typecheckers.
+# Does not exist at runtime.
+
+from collections.abc import Mapping, Sequence
+from typing import Any, Protocol
+
+DBAPITypeCode = Any | None
+# Strictly speaking, this should be a Sequence, but the type system does
+# not support fixed-length sequences.
+DBAPIColumnDescription = tuple[str, DBAPITypeCode, int | None, int | None, int | None, int | None, bool | None]
+
+class DBAPIConnection(Protocol):
+    def close(self) -> object: ...
+    def commit(self) -> object: ...
+    # optional:
+    # def rollback(self) -> Any: ...
+    def cursor(self) -> DBAPICursor: ...
+
+class DBAPICursor(Protocol):
+    @property
+    def description(self) -> Sequence[DBAPIColumnDescription] | None: ...
+    @property
+    def rowcount(self) -> int: ...
+    # optional:
+    # def callproc(self, __procname: str, __parameters: Sequence[Any] = ...) -> Sequence[Any]: ...
+    def close(self) -> object: ...
+    def execute(self, __operation: str, __parameters: Sequence[Any] | Mapping[str, Any] = ...) -> object: ...
+    def executemany(self, __operation: str, __seq_of_parameters: Sequence[Sequence[Any]]) -> object: ...
+    def fetchone(self) -> Sequence[Any] | None: ...
+    def fetchmany(self, __size: int = ...) -> Sequence[Sequence[Any]]: ...
+    def fetchall(self) -> Sequence[Sequence[Any]]: ...
+    # optional:
+    # def nextset(self) -> None | Literal[True]: ...
+    arraysize: int
+    def setinputsizes(self, __sizes: Sequence[DBAPITypeCode | int | None]) -> object: ...
+    def setoutputsize(self, __size: int, __column: int = ...) -> object: ...

stubs/SQLAlchemy/sqlalchemy/dialects/mysql/enumerated.pyi

@@ -3,7 +3,7 @@ from typing import Any
 from ...sql import sqltypes
 from .types import _StringType
 
-class ENUM(sqltypes.NativeForEmulated, sqltypes.Enum, _StringType):  # type: ignore[misc]
+class ENUM(sqltypes.NativeForEmulated, sqltypes.Enum, _StringType):  # type: ignore  # incompatible with base class
     __visit_name__: str
     native_enum: bool
     def __init__(self, *enums, **kw) -> None: ...

stubs/SQLAlchemy/sqlalchemy/dialects/mysql/types.pyi

@@ -32,7 +32,7 @@ class _StringType(sqltypes.String):
         **kw,
     ) -> None: ...
 
-class _MatchType(sqltypes.Float, sqltypes.MatchType):  # type: ignore[misc]
+class _MatchType(sqltypes.Float, sqltypes.MatchType):  # type: ignore  # incompatible with base class
     def __init__(self, **kw) -> None: ...
 
 class NUMERIC(_NumericType, sqltypes.NUMERIC):

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

@@ -96,7 +96,7 @@ PGUuid = UUID
 class TSVECTOR(sqltypes.TypeEngine):
     __visit_name__: str
 
-class ENUM(sqltypes.NativeForEmulated, sqltypes.Enum):  # type: ignore[misc]
+class ENUM(sqltypes.NativeForEmulated, sqltypes.Enum):  # type: ignore  # base classes incompatible
     native_enum: bool
     create_type: Any
     def __init__(self, *enums, **kw) -> None: ...