[SQLAlchemy] Annotate row classes (#9568)

Co-authored-by: Avasam <samuel.06@hotmail.com>

stubs/SQLAlchemy/@tests/stubtest_allowlist.txt

@@ -59,6 +59,12 @@ sqlalchemy.testing.provision.stop_test_class_outside_fixtures
 sqlalchemy.testing.provision.temp_table_keyword_args
 sqlalchemy.testing.provision.update_db_opts
 
+# potentially replaced at runtime
+sqlalchemy.engine.Row.count
+sqlalchemy.engine.Row.index
+sqlalchemy.engine.row.Row.count
+sqlalchemy.engine.row.Row.index
+
 # KeyError/AttributeError on import due to dynamic initialization from a different module
 sqlalchemy.testing.fixtures
 sqlalchemy.testing.pickleable

stubs/SQLAlchemy/sqlalchemy/cresultproxy.pyi

@@ -1,11 +1,23 @@
-from typing import Any
+from _typeshed import Incomplete
+from collections.abc import Callable, Iterable, Iterator
+from typing import Any, overload
 
 class BaseRow:
- def __init__(self, parent, processors, keymap, key_style, data) -> None: ...
- def __reduce__(self): ...
- def __iter__(self): ...
+ def __init__(
+ self,
+ __parent,
+ __processors: Iterable[Callable[[Any], Any]] | None,
+ __keymap: dict[Incomplete, Incomplete],
+ __key_style: int,
+ __row: Iterable[Any],
+ ) -> None: ...
+ def __reduce__(self) -> tuple[Incomplete, tuple[Incomplete, Incomplete]]: ...
+ def __iter__(self) -> Iterator[Any]: ...
  def __len__(self) -> int: ...
  def __hash__(self) -> int: ...
- __getitem__: Any
+ @overload
+ def __getitem__(self, __key: str | int) -> tuple[Any, ...]: ...
+ @overload
+ def __getitem__(self, __key: slice) -> tuple[tuple[Any, ...]]: ...
 
 def safe_rowproxy_reconstructor(__cls, __state): ...

stubs/SQLAlchemy/sqlalchemy/engine/row.pyi

@@ -1,9 +1,10 @@
-import abc
-from collections.abc import ItemsView, KeysView, Mapping, Sequence, ValuesView
-from typing import Any
+from collections.abc import ItemsView, Iterator, KeysView, Mapping, Sequence, ValuesView
+from typing import Any, Generic, TypeVar
 
 from ..cresultproxy import BaseRow as BaseRow
 
+_VT_co = TypeVar("_VT_co", covariant=True)
+
 MD_INDEX: int
 
 def rowproxy_reconstructor(cls, state): ...
@@ -13,45 +14,48 @@ KEY_OBJECTS_ONLY: int
 KEY_OBJECTS_BUT_WARN: int
 KEY_OBJECTS_NO_WARN: int
 
-class Row(BaseRow, Sequence[Any], metaclass=abc.ABCMeta):
+class Row(BaseRow, Sequence[Any]):
+ # The count and index methods are inherited from Sequence.
+ # If the result set contains columns with the same names, these
+ # fields contains their respective values, instead. We don't reflect
+ # this in the stubs.
+ __hash__ = BaseRow.__hash__ # type: ignore[assignment]
+ def __lt__(self, other: Row | tuple[Any, ...]) -> bool: ...
+ def __le__(self, other: Row | tuple[Any, ...]) -> bool: ...
+ def __ge__(self, other: Row | tuple[Any, ...]) -> bool: ...
+ def __gt__(self, other: Row | tuple[Any, ...]) -> bool: ...
+ def __eq__(self, other: object) -> bool: ...
+ def __ne__(self, other: object) -> bool: ...
+ def keys(self) -> list[str]: ...
+ # The following methods are public, but have a leading underscore
+ # to prevent conflicts with column names.
  @property
- def count(self): ...
+ def _mapping(self) -> RowMapping: ...
  @property
- def index(self): ...
- def __contains__(self, key): ...
- __hash__ = BaseRow.__hash__
- def __lt__(self, other): ...
- def __le__(self, other): ...
- def __ge__(self, other): ...
- def __gt__(self, other): ...
- def __eq__(self, other): ...
- def __ne__(self, other): ...
- def keys(self): ...
-
-class LegacyRow(Row, metaclass=abc.ABCMeta):
- def __contains__(self, key): ...
- def has_key(self, key): ...
- def items(self): ...
- def iterkeys(self): ...
- def itervalues(self): ...
- def values(self): ...
+ def _fields(self) -> tuple[str, ...]: ...
+ def _asdict(self) -> dict[str, Any]: ...
+
+class LegacyRow(Row):
+ def has_key(self, key: str) -> bool: ...
+ def items(self) -> list[tuple[str, Any]]: ...
+ def iterkeys(self) -> Iterator[str]: ...
+ def itervalues(self) -> Iterator[Any]: ...
+ def values(self) -> list[Any]: ...
 
 BaseRowProxy = BaseRow
 RowProxy = Row
 
-class ROMappingView(KeysView[Any], ValuesView[Any], ItemsView[Any, Any]):
- def __init__(self, mapping, items) -> None: ...
+class ROMappingView(KeysView[str], ValuesView[_VT_co], ItemsView[str, _VT_co], Generic[_VT_co]): # type: ignore[misc]
+ def __init__(self, mapping: RowMapping, items: list[_VT_co]) -> None: ...
  def __len__(self) -> int: ...
- def __iter__(self): ...
- def __contains__(self, item): ...
- def __eq__(self, other): ...
- def __ne__(self, other): ...
+ def __iter__(self) -> Iterator[_VT_co]: ... # type: ignore[override]
+ def __eq__(self, other: ROMappingView[_VT_co]) -> bool: ... # type: ignore[override]
+ def __ne__(self, other: ROMappingView[_VT_co]) -> bool: ... # type: ignore[override]
 
-class RowMapping(BaseRow, Mapping[Any, Any]):
+class RowMapping(BaseRow, Mapping[str, Row]):
  __getitem__: Any
- def __iter__(self): ...
+ def __iter__(self) -> Iterator[str]: ...
  def __len__(self) -> int: ...
- def __contains__(self, key): ...
- def items(self): ...
- def keys(self): ...
- def values(self): ...
+ def items(self) -> ROMappingView[tuple[str, Any]]: ... # type: ignore[override]
+ def keys(self) -> list[str]: ... # type: ignore[override]
+ def values(self) -> ROMappingView[Any]: ... # type: ignore[override]