diff --git a/crates/ty_python_semantic/resources/mdtest/external/sqlalchemy.md b/crates/ty_python_semantic/resources/mdtest/external/sqlalchemy.md index 0ac9c4c219..d4a01c7701 100644 --- a/crates/ty_python_semantic/resources/mdtest/external/sqlalchemy.md +++ b/crates/ty_python_semantic/resources/mdtest/external/sqlalchemy.md @@ -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 +```