[ty] Increase SQLAlchemy test coverage (#21843)

## Summary

Increase our SQLAlchemy test coverage to make sure we understand
`Session.scalar`, `Session.scalars`, `Session.execute` (and their async
equivalents), as well as `Result.tuples`, `Result.one_or_none`,
`Row._tuple`.
This commit is contained in:
David Peter 2025-12-08 14:36:13 +01:00 committed by GitHub
parent a364195335
commit c99e10eedc
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 103 additions and 19 deletions

View File

@ -9,9 +9,9 @@ python-platform = "linux"
dependencies = ["SQLAlchemy==2.0.44"]
```
## Basic model
## ORM Model
Here, we mostly make sure that ty understands SQLAlchemy's dataclass-transformer setup:
This test makes sure that ty understands SQLAlchemy's `dataclass_transform` setup:
```py
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
@ -40,9 +40,9 @@ reveal_type(User.__init__) # revealed: def __init__(self, **kw: Any) -> Unknown
invalid_user = User(invalid_arg=42)
```
## Queries
## Basic query example
First, the basic setup:
First, set up a `Session`:
```py
from datetime import datetime
@ -57,7 +57,7 @@ engine = create_engine("sqlite://example.db")
session = Session(engine)
```
Now we can declare a simple model:
And define a simple model:
```py
class Base(DeclarativeBase):
@ -71,7 +71,7 @@ class User(Base):
is_admin: Mapped[bool] = mapped_column(Boolean, default=False)
```
And perform simple queries:
Finally, we can execute queries:
```py
stmt = select(User)
@ -84,21 +84,27 @@ for row in session.execute(stmt):
reveal_type(row) # revealed: Row[tuple[User]]
stmt = select(User).where(User.name == "Alice")
alice = session.scalars(stmt).first()
reveal_type(alice) # revealed: User | None
alice1 = session.scalars(stmt).first()
reveal_type(alice1) # revealed: User | None
alice2 = session.scalar(stmt)
reveal_type(alice2) # revealed: User | None
result = session.execute(stmt)
row = result.one_or_none()
assert row is not None
(alice3,) = row._tuple()
reveal_type(alice3) # revealed: User
```
This also works with more complex queries:
```py
stmt = select(User).where(User.is_admin == True).order_by(User.name).limit(10)
admin_users = session.scalars(stmt).all()
reveal_type(admin_users) # revealed: Sequence[User]
```
This also works with the legacy `query` API:
```py
users_legacy = session.query(User).all()
reveal_type(users_legacy) # revealed: list[User]
```
We can also specify particular columns to select:
```py
@ -106,19 +112,97 @@ stmt = select(User.id, User.name)
# TODO: should be `Select[tuple[int, str]]`
reveal_type(stmt) # revealed: Select[tuple[Unknown, Unknown]]
ids_and_names = session.execute(stmt).all()
# TODO: should be `Sequence[Row[tuple[int, str]]]`
reveal_type(ids_and_names) # revealed: Sequence[Row[tuple[Unknown, Unknown]]]
for row in session.execute(stmt):
# TODO: should be `Row[Tuple[int, str]]`
# TODO: should be `Row[tuple[int, str]]`
reveal_type(row) # revealed: Row[tuple[Unknown, Unknown]]
for user_id, name in session.execute(stmt).tuples():
# TODO: should be `int`
reveal_type(user_id) # revealed: Unknown
# TODO: should be `str`
reveal_type(name) # revealed: Unknown
stmt = select(User.id, User.name).where(User.name == "Alice")
alice1 = session.scalars(stmt).first()
# TODO: should be `tuple[int, str] | None`
reveal_type(alice1) # revealed: Any | None
alice2 = session.scalar(stmt)
# TODO: should be `tuple[int, str] | None`
reveal_type(alice2) # revealed: Any
result = session.execute(stmt)
row = result.one_or_none()
assert row is not None
(user_id, name) = row._tuple()
# TODO: should be `int`
reveal_type(user_id) # revealed: Unknown
# TODO: should be `str`
reveal_type(name) # revealed: Unknown
```
And similarly with the legacy `query` API:
Using the legacy `query` API also works:
```py
users_legacy = session.query(User).all()
reveal_type(users_legacy) # revealed: list[User]
query = session.query(User)
reveal_type(query) # revealed: Query[User]
reveal_type(query.all()) # revealed: list[User]
for row in query:
reveal_type(row) # revealed: User
```
And similarly when specifying particular columns:
```py
query = session.query(User.id, User.name)
# TODO: should be `RowReturningQuery[tuple[int, str]]`
reveal_type(query) # revealed: RowReturningQuery[tuple[Unknown, Unknown]]
for row in query.all():
# TODO: should be `Row[Tuple[int, str]]`
# TODO: should be `list[Row[tuple[int, str]]]`
reveal_type(query.all()) # revealed: list[Row[tuple[Unknown, Unknown]]]
for row in query:
# TODO: should be `Row[tuple[int, str]]`
reveal_type(row) # revealed: Row[tuple[Unknown, Unknown]]
```
## Async API
The async API is supported as well:
```py
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select, Integer, Text
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
class Base(DeclarativeBase):
pass
class User(Base):
__tablename__ = "users"
id: Mapped[int] = mapped_column(Integer, primary_key=True)
name: Mapped[str] = mapped_column(Text)
async def test_async(session: AsyncSession):
stmt = select(User).where(User.name == "Alice")
alice = await session.scalar(stmt)
reveal_type(alice) # revealed: User | None
stmt = select(User.id, User.name)
result = await session.execute(stmt)
for user_id, name in result.tuples():
# TODO: should be `int`
reveal_type(user_id) # revealed: Unknown
# TODO: should be `str`
reveal_type(name) # revealed: Unknown
```