mirror of https://github.com/astral-sh/ruff
[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:
parent
a364195335
commit
c99e10eedc
|
|
@ -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
|
||||
```
|
||||
|
|
|
|||
Loading…
Reference in New Issue