python · Avasam · Jan 27, 2023 · Jan 27, 2023 · Jan 28, 2023 · Jan 29, 2023
diff --git a/stubs/SQLAlchemy/@tests/stubtest_allowlist.txt b/stubs/SQLAlchemy/@tests/stubtest_allowlist.txt
@@ -5,12 +5,17 @@ sqlalchemy.testing.suite
 # Leaked re-exports from the compat module
 sqlalchemy.util.quote
 
+# Parameter's default value would always fail
+sqlalchemy.sql.elements.Slice.self_group
 # Expanding keyword arguments in stubs
 sqlalchemy.ext.declarative.as_declarative
 
 # not always present
-sqlalchemy.engine.Engine.logging_name  # initialized if not None
+sqlalchemy.cresultproxy.BaseRow.__getattr__
+sqlalchemy.engine.BaseRow.__getattr__
+sqlalchemy.engine.row.BaseRow.__getattr__
 sqlalchemy.engine.base.Engine.logging_name  # initialized if not None
+sqlalchemy.engine.Engine.logging_name  # initialized if not None
 sqlalchemy.sql.lambdas.PyWrapper.__clause_element__
 sqlalchemy.testing.util.non_refcount_gc_collect
 
@@ -25,23 +30,58 @@ sqlalchemy.engine.Transaction.connection
 sqlalchemy.engine.Transaction.is_active
 sqlalchemy.engine.base.Transaction.connection
 sqlalchemy.engine.base.Transaction.is_active
+# @type_check_only
+sqlalchemy.orm.RelationshipProperty.__get__
+sqlalchemy.orm.properties.RelationshipProperty.__get__
+sqlalchemy.orm.relationships.RelationshipProperty.__get__
 
 # initialized to None during class construction, but overridden during __init__() or __new__()
+sqlalchemy.Column.foreign_keys
+sqlalchemy.Column.key
 sqlalchemy.engine.base.Connection.engine
+sqlalchemy.engine.base.ExceptionContextImpl.is_disconnect
 sqlalchemy.engine.Connection.engine
-sqlalchemy.orm.Mapper.single
+sqlalchemy.ForeignKeyConstraint.columns
+sqlalchemy.ForeignKeyConstraint.elements
+sqlalchemy.MetaData.tables
+sqlalchemy.orm.Mapper.concrete
+sqlalchemy.orm.mapper.Mapper.concrete
+sqlalchemy.orm.mapper.Mapper.non_primary
 sqlalchemy.orm.mapper.Mapper.single
-
-# Uses @memoized_property at runtime, but we use @property for compatibility
-sqlalchemy.engine.URL.normalized_query
-sqlalchemy.engine.url.URL.normalized_query
-# Uses @memoized_property, but that causes regr_test to raise 'Cannot determine type of "..." in base class "..."  [misc]'
-sqlalchemy.schema.SchemaItem.info
-sqlalchemy.sql.elements.AnnotatedColumnElement.info
-sqlalchemy.sql.elements.AnnotatedColumnElement.key
-sqlalchemy.sql.elements.AnnotatedColumnElement.name
-sqlalchemy.sql.elements.AnnotatedColumnElement.table
-sqlalchemy.sql.schema.SchemaItem.info
+sqlalchemy.orm.Mapper.non_primary
+sqlalchemy.orm.Mapper.single
+sqlalchemy.schema.Column.foreign_keys
+sqlalchemy.schema.Column.key
+sqlalchemy.schema.ColumnCollectionConstraint.columns
+sqlalchemy.schema.ColumnCollectionMixin.columns
+sqlalchemy.schema.ForeignKeyConstraint.columns
+sqlalchemy.schema.ForeignKeyConstraint.elements
+sqlalchemy.schema.MetaData.tables
+sqlalchemy.schema.Table.constraints
+sqlalchemy.schema.Table.indexes
+sqlalchemy.sql.elements.BindParameter.key
+sqlalchemy.sql.elements.ColumnClause.key
+sqlalchemy.sql.elements.Label.key
+sqlalchemy.sql.elements.Over.element
+sqlalchemy.sql.elements.TableValuedColumn.key
+sqlalchemy.sql.expression.BindParameter.key
+sqlalchemy.sql.expression.ColumnClause.key
+sqlalchemy.sql.expression.Label.key
+sqlalchemy.sql.expression.Over.element
+sqlalchemy.sql.schema.Column.foreign_keys
+sqlalchemy.sql.schema.Column.key
+sqlalchemy.sql.schema.ColumnCollectionConstraint.columns
+sqlalchemy.sql.schema.ColumnCollectionMixin.columns
+sqlalchemy.sql.schema.ForeignKeyConstraint.columns
+sqlalchemy.sql.schema.ForeignKeyConstraint.elements
+sqlalchemy.sql.schema.MetaData.tables
+sqlalchemy.sql.schema.Table.constraints
+sqlalchemy.sql.schema.Table.indexes
+sqlalchemy.sql.selectable.BindParameter.key
+sqlalchemy.sql.selectable.ColumnClause.key
+sqlalchemy.sql.selectable.TableValuedColumn.key
+sqlalchemy.Table.constraints
+sqlalchemy.Table.indexes
 
 # runtime has extra internal arguments that are inconsistent across micro versions
 sqlalchemy.testing.engines.testing_engine
@@ -60,6 +100,9 @@ sqlalchemy.testing.fixtures
 sqlalchemy.testing.pickleable
 sqlalchemy.testing.plugin.bootstrap
 
+# unclear problems
+sqlalchemy.sql.functions.array_agg.type
+
 # method arguments starting with double underscores in the implementation trips up stubtest
 sqlalchemy.testing.resolve_lambda
 sqlalchemy.testing.util.resolve_lambda

diff --git a/stubs/SQLAlchemy/sqlalchemy/__init__.pyi b/stubs/SQLAlchemy/sqlalchemy/__init__.pyi
@@ -1,3 +1,4 @@
+from . import util
 from .engine import (
     create_engine as create_engine,
     create_mock_engine as create_mock_engine,
@@ -130,4 +131,5 @@ from .sql.sqltypes import (
     UnicodeText as UnicodeText,
 )
 
+_util = util
 __version__: str
diff --git a/stubs/SQLAlchemy/sqlalchemy/connectors/mxodbc.pyi b/stubs/SQLAlchemy/sqlalchemy/connectors/mxodbc.pyi
@@ -12,6 +12,6 @@ class MxODBCConnector(Connector):
     def dbapi(cls): ...
     def on_connect(self): ...
     def create_connect_args(self, url): ...
-    def is_disconnect(self, e, connection, cursor): ...
+    def is_disconnect(self, e, connection, cursor) -> bool: ...
     def do_executemany(self, cursor, statement, parameters, context: Incomplete | None = None) -> None: ...
     def do_execute(self, cursor, statement, parameters, context: Incomplete | None = None) -> None: ...
diff --git a/stubs/SQLAlchemy/sqlalchemy/connectors/pyodbc.pyi b/stubs/SQLAlchemy/sqlalchemy/connectors/pyodbc.pyi
@@ -17,6 +17,6 @@ class PyODBCConnector(Connector):
     @classmethod
     def dbapi(cls): ...
     def create_connect_args(self, url): ...
-    def is_disconnect(self, e, connection, cursor): ...
+    def is_disconnect(self, e, connection, cursor) -> bool: ...
     def do_set_input_sizes(self, cursor, list_of_tuples, context) -> None: ...
     def set_isolation_level(self, connection, level) -> None: ...
diff --git a/stubs/SQLAlchemy/sqlalchemy/cresultproxy.pyi b/stubs/SQLAlchemy/sqlalchemy/cresultproxy.pyi
@@ -1,23 +1,96 @@
 from _typeshed import Incomplete
-from collections.abc import Callable, Iterable, Iterator
-from typing import Any, overload
+from collections.abc import Callable, Iterable, Iterator, Sequence
+from typing import TypeVar, overload
+from typing_extensions import Literal
 
-class BaseRow:
+from .engine.result import ResultMetaData
+
+_T = TypeVar("_T")
+_T_co = TypeVar("_T_co")
+# _KS = TypeVar("_KS", Literal[0], Literal[1], Literal[2, 3], int)
+
+class BaseRow(Iterable[_T_co]):
+    @overload
     def __init__(
         self,
-        __parent,
-        __processors: Iterable[Callable[[Any], Any]] | None,
+        __parent: ResultMetaData,
+        __processors: Iterable[Callable[[_T_co], Incomplete] | None] | None,
+        __keymap: dict[Incomplete, Incomplete],
+        __key_style: Literal[1],
+        __row: Sequence[_T_co],
+    ) -> None: ...
+    @overload
+    def __init__(
+        self,
+        __parent: ResultMetaData | None,
+        __processors: Iterable[Callable[[_T_co], Incomplete]] | None,
         __keymap: dict[Incomplete, Incomplete],
         __key_style: int,
-        __row: Iterable[Any],
+        __row: Sequence[_T_co],
     ) -> None: ...
     def __reduce__(self) -> tuple[Incomplete, tuple[Incomplete, Incomplete]]: ...
-    def __iter__(self) -> Iterator[Any]: ...
+    def __iter__(self) -> Iterator[_T_co]: ...
     def __len__(self) -> int: ...
     def __hash__(self) -> int: ...
+    def __getattr__(self, name: str) -> _T_co: ...
     @overload
-    def __getitem__(self, __key: str | int) -> tuple[Any, ...]: ...
+    def __getitem__(self, index: slice) -> tuple[_T_co, ...]: ...  # type: ignore[misc]
     @overload
-    def __getitem__(self, __key: slice) -> tuple[tuple[Any, ...]]: ...
+    def __getitem__(self, index: object) -> _T_co: ...
+
+# TODO: Keeping BaseRow non-Generic for now to reduce the amount of changes seen by mypy_primer
+# Change to the implementation below in a different PR
+# class BaseRow(Iterable[_T_co], Generic[_T_co, _KS]):
+#     @overload
+#     def __init__(
+#         self
+#         __parent: ResultMetaData,
+#         __processors: Iterable[Callable[[_T_co], Incomplete]] | None,
+#         __keymap: dict[Incomplete, Incomplete],
+#         __key_style: Literal[1],
+#         __row: Sequence[_T_co],
+#     ) -> None: ...
+#     @overload
+#     def __init__(
+#         self
+#         __parent: ResultMetaData | None,
+#         __processors: Iterable[Callable[[_T_co], Incomplete]] | None,
+#         __keymap: dict[Incomplete, Incomplete],
+#         __key_style: _KS,
+#         __row: Sequence[_T_co],
+#     ) -> None: ...
+#     def __reduce__(self) -> tuple[Incomplete, tuple[Incomplete, Incomplete]]: ...
+#     def __iter__(self) -> Iterator[_T_co]: ...
+#     def __len__(self) -> int: ...
+#     def __hash__(self) -> int: ...
+#     def __getattr__(self, name: str) -> _T_co: ...
+#     ###
+#     # __key_style = 0
+#     ###
+#     @overload
+#     def __getitem__(self: BaseRow[_T_co, Literal[0]], index: int) -> _T_co: ...
+#     @overload
+#     def __getitem__(self: BaseRow[_T_co, Literal[0]], index: slice) -> tuple[_T_co, ...]: ...  # type: ignore[misc]
+#     # Needed because the int overload would assume index=object|str is aceptable with __key_style=0
+#     @overload
+#     def __getitem__(  # type: ignore[misc]
+#         self: BaseRow[_T_co, Literal[0]], index: object
+#     ) -> _T_co | tuple[_T_co, ...] | NoReturn: ...
+#     ###
+#     # __key_style = 1
+#     ###
+#     # Needed because next overload is "object & Not[int | slice]"
+#     @overload
+#     def __getitem__(self: BaseRow[_T_co, Literal[1]], index: int | slice) -> NoReturn: ...
+#     @overload
+#     def __getitem__(self: BaseRow[_T_co, Literal[1]], index: object) -> _T_co: ...  # type: ignore[misc]
+#     ###
+#     # __key_style = 3 | 4
+#     ###
+#     # General fallbacks if _KS evaluates to int
+#     @overload
+#     def __getitem__(self: BaseRow[_T_co, int], index: slice) -> tuple[_T_co, ...]: ...  # type: ignore[misc]
+#     @overload
+#     def __getitem__(self: BaseRow[_T_co, int], index: object) -> _T_co: ...
 
-def safe_rowproxy_reconstructor(__cls, __state): ...
+def safe_rowproxy_reconstructor(__cls: type[_T], __state: dict[str, Incomplete]) -> _T: ...
diff --git a/stubs/SQLAlchemy/sqlalchemy/dialects/__init__.pyi b/stubs/SQLAlchemy/sqlalchemy/dialects/__init__.pyi
@@ -1,5 +1,4 @@
-from typing import Any
-
+from ..util.langhelpers import PluginLoader
 from . import (
     firebird as firebird,
     mssql as mssql,
@@ -12,5 +11,5 @@ from . import (
 
 __all__ = ("firebird", "mssql", "mysql", "oracle", "postgresql", "sqlite", "sybase")
 
-registry: Any
-plugins: Any
+registry: PluginLoader
+plugins: PluginLoader
diff --git a/stubs/SQLAlchemy/sqlalchemy/dialects/firebird/__init__.pyi b/stubs/SQLAlchemy/sqlalchemy/dialects/firebird/__init__.pyi
@@ -1,5 +1,4 @@
-from typing import Any
-
+from . import base as base, fdb as fdb, kinterbasdb as kinterbasdb
 from .base import (
     BIGINT as BIGINT,
     BLOB as BLOB,
@@ -31,4 +30,4 @@ __all__ = (
     "dialect",
 )
 
-dialect: Any
+dialect = fdb.dialect
diff --git a/stubs/SQLAlchemy/sqlalchemy/dialects/firebird/base.pyi b/stubs/SQLAlchemy/sqlalchemy/dialects/firebird/base.pyi
@@ -96,8 +96,8 @@ class FBDialect(default.DefaultDialect):
     def __init__(self, *args, **kwargs) -> None: ...
     implicit_returning: Any
     def initialize(self, connection) -> None: ...
-    def has_table(self, connection, table_name, schema: Incomplete | None = None): ...  # type: ignore[override]
-    def has_sequence(self, connection, sequence_name, schema: Incomplete | None = None): ...  # type: ignore[override]
+    def has_table(self, connection, table_name, schema: Incomplete | None = None) -> bool: ...  # type: ignore[override]
+    def has_sequence(self, connection, sequence_name, schema: Incomplete | None = None) -> bool: ...  # type: ignore[override]
     def get_table_names(self, connection, schema: Incomplete | None = None, **kw): ...
     def get_view_names(self, connection, schema: Incomplete | None = None, **kw): ...
     def get_view_definition(self, connection, view_name, schema: Incomplete | None = None, **kw): ...

diff --git a/stubs/SQLAlchemy/sqlalchemy/dialects/firebird/kinterbasdb.pyi b/stubs/SQLAlchemy/sqlalchemy/dialects/firebird/kinterbasdb.pyi
@@ -34,6 +34,6 @@ class FBDialect_kinterbasdb(FBDialect):
     def do_rollback(self, dbapi_connection) -> None: ...
     def do_commit(self, dbapi_connection) -> None: ...
     def create_connect_args(self, url): ...
-    def is_disconnect(self, e, connection, cursor): ...
+    def is_disconnect(self, e, connection, cursor) -> bool: ...
 
 dialect = FBDialect_kinterbasdb
diff --git a/stubs/SQLAlchemy/sqlalchemy/dialects/mssql/__init__.pyi b/stubs/SQLAlchemy/sqlalchemy/dialects/mssql/__init__.pyi
@@ -1,5 +1,4 @@
-from typing import Any
-
+from . import base as base, mxodbc as mxodbc, pymssql as pymssql, pyodbc as pyodbc
 from .base import (
     BIGINT as BIGINT,
     BINARY as BINARY,
@@ -73,4 +72,4 @@ __all__ = (
     "try_cast",
 )
 
-dialect: Any
+dialect = pyodbc.dialect
diff --git a/stubs/SQLAlchemy/sqlalchemy/dialects/mssql/base.pyi b/stubs/SQLAlchemy/sqlalchemy/dialects/mssql/base.pyi
@@ -1,4 +1,5 @@
 from _typeshed import Incomplete
+from abc import ABCMeta
 from typing import Any, overload
 from typing_extensions import Literal
 
@@ -81,7 +82,7 @@ class _MSUnicodeText(_UnicodeLiteral, sqltypes.UnicodeText): ...
 class TIMESTAMP(sqltypes._Binary):
     __visit_name__: str
     length: Any
-    convert_int: Any
+    convert_int: bool
     def __init__(self, convert_int: bool = False) -> None: ...
     def result_processor(self, dialect, coltype): ...
 
@@ -240,7 +241,7 @@ class MSSQLCompiler(compiler.SQLCompiler):
     def visit_json_path_getitem_op_binary(self, binary, operator, **kw): ...
     def visit_sequence(self, seq, **kw): ...
 
-class MSSQLStrictCompiler(MSSQLCompiler):
+class MSSQLStrictCompiler(MSSQLCompiler, metaclass=ABCMeta):
     ansi_bind_rules: bool
     def visit_in_op_binary(self, binary, operator, **kw): ...
     def visit_not_in_op_binary(self, binary, operator, **kw): ...
@@ -310,8 +311,8 @@ class MSDialect(default.DefaultDialect):
     def get_isolation_level(self, dbapi_connection): ...
     def initialize(self, connection) -> None: ...
     def on_connect(self): ...
-    def has_table(self, connection, tablename, dbname, owner, schema): ...
-    def has_sequence(self, connection, sequencename, dbname, owner, schema): ...
+    def has_table(self, connection, tablename, dbname, owner, schema) -> bool: ...  # type: ignore[override]  # Actual bool
+    def has_sequence(self, connection, sequencename, dbname, owner, schema) -> bool: ...  # type: ignore[override]  # Actual bool
     def get_sequence_names(self, connection, dbname, owner, schema, **kw): ...
     def get_schema_names(self, connection, **kw): ...
     def get_table_names(self, connection, dbname, owner, schema, **kw): ...

diff --git a/stubs/SQLAlchemy/sqlalchemy/dialects/mssql/json.pyi b/stubs/SQLAlchemy/sqlalchemy/dialects/mssql/json.pyi
@@ -1,10 +1,10 @@
-from ...sql.sqltypes import JSON as _JSON
+from ...sql import sqltypes
 
-class JSON(_JSON): ...
+class JSON(sqltypes.JSON): ...
 
 class _FormatTypeMixin:
     def bind_processor(self, dialect): ...
     def literal_processor(self, dialect): ...
 
-class JSONIndexType(_FormatTypeMixin, _JSON.JSONIndexType): ...
-class JSONPathType(_FormatTypeMixin, _JSON.JSONPathType): ...
+class JSONIndexType(_FormatTypeMixin, sqltypes.JSON.JSONIndexType): ...
+class JSONPathType(_FormatTypeMixin, sqltypes.JSON.JSONPathType): ...
diff --git a/stubs/SQLAlchemy/sqlalchemy/dialects/mssql/pymssql.pyi b/stubs/SQLAlchemy/sqlalchemy/dialects/mssql/pymssql.pyi
@@ -18,7 +18,7 @@ class MSDialect_pymssql(MSDialect):
     @classmethod
     def dbapi(cls): ...
     def create_connect_args(self, url): ...
-    def is_disconnect(self, e, connection, cursor): ...
+    def is_disconnect(self, e, connection, cursor) -> bool: ...
     def set_isolation_level(self, connection, level) -> None: ...
 
 dialect = MSDialect_pymssql
diff --git a/stubs/SQLAlchemy/sqlalchemy/dialects/mssql/pyodbc.pyi b/stubs/SQLAlchemy/sqlalchemy/dialects/mssql/pyodbc.pyi
@@ -40,6 +40,6 @@ class MSDialect_pyodbc(PyODBCConnector, MSDialect):
     def __init__(self, description_encoding: Incomplete | None = None, fast_executemany: bool = False, **params) -> None: ...
     def on_connect(self): ...
     def do_executemany(self, cursor, statement, parameters, context: Incomplete | None = None) -> None: ...
-    def is_disconnect(self, e, connection, cursor): ...
+    def is_disconnect(self, e, connection, cursor) -> bool: ...
 
 dialect = MSDialect_pyodbc
diff --git a/stubs/SQLAlchemy/sqlalchemy/dialects/mysql/__init__.pyi b/stubs/SQLAlchemy/sqlalchemy/dialects/mysql/__init__.pyi
@@ -1,5 +1,15 @@
-from typing import Any
-
+from . import (
+    aiomysql as aiomysql,
+    asyncmy as asyncmy,
+    base as base,
+    cymysql as cymysql,
+    mariadbconnector as mariadbconnector,
+    mysqlconnector as mysqlconnector,
+    mysqldb as mysqldb,
+    oursql as oursql,
+    pymysql as pymysql,
+    pyodbc as pyodbc,
+)
 from .base import (
     BIGINT as BIGINT,
     BINARY as BINARY,
@@ -82,4 +92,4 @@ __all__ = (
     "match",
 )
 
-dialect: Any
+dialect = mysqldb.dialect
diff --git a/stubs/SQLAlchemy/sqlalchemy/dialects/mysql/aiomysql.pyi b/stubs/SQLAlchemy/sqlalchemy/dialects/mysql/aiomysql.pyi
@@ -68,7 +68,7 @@ class MySQLDialect_aiomysql(MySQLDialect_pymysql):
     @classmethod
     def get_pool_class(cls, url): ...
     def create_connect_args(self, url): ...
-    def is_disconnect(self, e, connection, cursor): ...
+    def is_disconnect(self, e, connection, cursor) -> bool: ...
     def get_driver_connection(self, connection): ...
 
 dialect = MySQLDialect_aiomysql