Compare commits

..

6 Commits

Author SHA1 Message Date
David Peter
ad4945aed0 SQLAlchemy investigation 2025-12-08 15:30:58 +01:00
David Peter
4686111681 [ty] More SQLAlchemy test updates (#21846)
Minor updates to the SQLAlchemy test suite. I verified all expected
results using pyright.
2025-12-08 15:22:55 +01:00
Micha Reiser
4364ffbdd3 [ty] Don't create a related diagnostic for the primary annotation of sub-diagnostics (#21845) 2025-12-08 14:22:11 +00:00
Charlie Marsh
b845e81c4a Use memchr for computing line indexes (#21838)
## Summary

Some benchmarks with Claude's help:

| File | Size | Baseline | Optimized | Speedup |

|---------------------|-------|----------------------|----------------------|---------|
| numpy/globals.py | 3 KB | 1.48 µs (1.95 GiB/s) | 740 ns (3.89 GiB/s) |
2.0x |
| unicode/pypinyin.py | 4 KB | 2.04 µs (2.01 GiB/s) | 1.18 µs (3.49
GiB/s) | 1.7x |
| pydantic/types.py | 26 KB | 13.1 µs (1.90 GiB/s) | 5.88 µs (4.23
GiB/s) | 2.2x |
| numpy/ctypeslib.py | 17 KB | 8.45 µs (1.92 GiB/s) | 3.94 µs (4.13
GiB/s) | 2.1x |
| large/dataset.py | 41 KB | 21.6 µs (1.84 GiB/s) | 11.2 µs (3.55 GiB/s)
| 1.9x |

I think that I originally thought we _had_ to iterate
character-by-character here because we needed to do the ASCII check, but
the ASCII check can be vectorized by LLVM (and the "search for newlines"
can be done with `memchr`).
2025-12-08 08:50:51 -05:00
David Peter
c99e10eedc [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`.
2025-12-08 14:36:13 +01:00
Dhruv Manilawala
a364195335 [ty] Avoid diagnostic when typing_extensions.ParamSpec uses default parameter (#21839)
## Summary

fixes: https://github.com/astral-sh/ty/issues/1798

## Test Plan

Add mdtest.
2025-12-08 12:34:30 +00:00
9 changed files with 366 additions and 64 deletions

View File

@@ -888,6 +888,10 @@ impl Annotation {
pub fn hide_snippet(&mut self, yes: bool) {
self.hide_snippet = yes;
}
pub fn is_primary(&self) -> bool {
self.is_primary
}
}
/// Tags that can be associated with an annotation.

View File

@@ -33,26 +33,29 @@ impl LineIndex {
line_starts.push(TextSize::default());
let bytes = text.as_bytes();
let mut utf8 = false;
assert!(u32::try_from(bytes.len()).is_ok());
for (i, byte) in bytes.iter().enumerate() {
utf8 |= !byte.is_ascii();
match byte {
// Only track one line break for `\r\n`.
b'\r' if bytes.get(i + 1) == Some(&b'\n') => continue,
b'\n' | b'\r' => {
// SAFETY: Assertion above guarantees `i <= u32::MAX`
#[expect(clippy::cast_possible_truncation)]
line_starts.push(TextSize::from(i as u32) + TextSize::from(1));
}
_ => {}
for i in memchr::memchr2_iter(b'\n', b'\r', bytes) {
// Skip `\r` in `\r\n` sequences (only count the `\n`).
if bytes[i] == b'\r' && bytes.get(i + 1) == Some(&b'\n') {
continue;
}
// SAFETY: Assertion above guarantees `i <= u32::MAX`
#[expect(clippy::cast_possible_truncation)]
line_starts.push(TextSize::from(i as u32) + TextSize::from(1));
}
let kind = if utf8 {
// Determine whether the source text is ASCII.
//
// Empirically, this simple loop is auto-vectorized by LLVM and benchmarks faster than both
// `str::is_ascii()` and hand-written SIMD.
let mut has_non_ascii = false;
for byte in bytes {
has_non_ascii |= !byte.is_ascii();
}
let kind = if has_non_ascii {
IndexKind::Utf8
} else {
IndexKind::Ascii

View File

@@ -82,24 +82,6 @@ async def main():
reveal_type(b) # revealed: int
```
### `asynccontextmanager`
```py
from contextlib import asynccontextmanager
from typing import AsyncGenerator
class Session: ...
@asynccontextmanager
async def connect() -> AsyncGenerator[Session]:
yield Session()
async def main():
async with connect() as session:
# TODO: should be `Session`
reveal_type(session) # revealed: Unknown
```
## Under the hood
```toml

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,14 +40,12 @@ 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
from sqlalchemy import select, Integer, Text, Boolean, DateTime
from sqlalchemy import select, Integer, Text, Boolean
from sqlalchemy.orm import Session
from sqlalchemy.orm import DeclarativeBase
from sqlalchemy.orm import Mapped, mapped_column
@@ -57,7 +55,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 +69,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 +82,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 +110,245 @@ 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
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
stmt = select(User.id).where(User.name == "Alice")
# TODO: should be `Select[tuple[int]]`
reveal_type(stmt) # revealed: Select[tuple[Unknown]]
alice_id = session.scalars(stmt).first()
# TODO: should be `int | None`
reveal_type(alice_id) # revealed: Unknown | None
alice_id = session.scalar(stmt)
# TODO: should be `int | None`
reveal_type(alice_id) # revealed: Unknown | None
```
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
```
## What is it that we do not support yet?
Basic setup:
```py
from datetime import datetime
from sqlalchemy import select, Integer, Text, Boolean, DateTime
from sqlalchemy.orm import Session
from sqlalchemy.orm import DeclarativeBase
from sqlalchemy.orm import Mapped, mapped_column
from sqlalchemy import create_engine
engine = create_engine("sqlite://example.db")
session = Session(engine)
class Base(DeclarativeBase):
pass
class User(Base):
__tablename__ = "users"
id: Mapped[int] = mapped_column(Integer, primary_key=True)
name: Mapped[str] = mapped_column(Text)
is_admin: Mapped[bool] = mapped_column(Boolean, default=False)
```
Why do we see `Unknown`s for `select(User.id, User.name)` here?
```py
stmt = select(User.id, User.name)
# TODO: should be `Select[tuple[int, str]]`
reveal_type(stmt) # revealed: Select[tuple[Unknown, Unknown]]
```
The types of the arguments seem correct:
```py
reveal_type(User.id) # revealed: InstrumentedAttribute[int]
reveal_type(User.name) # revealed: InstrumentedAttribute[str]
```
The two-parameter overload of `select` has a type of
`def select(__ent0: _TCCA[_T0], __ent1: _TCCA[_T1], /) -> Select[_T0, _T1]: ...`
here `_TCCA` is an alias for `_TypedColumnClauseArgument`:
```py
from sqlalchemy.sql._typing import _TypedColumnClauseArgument
# revealed: <types.UnionType special form 'TypedColumnsClauseRole[_T@_TypedColumnClauseArgument] | SQLCoreOperations[_T@_TypedColumnClauseArgument] | type[_T@_TypedColumnClauseArgument]'>
reveal_type(_TypedColumnClauseArgument)
```
If we use that generic type alias in a type expression, we can properly specialize it:
```py
def _(
col: _TypedColumnClauseArgument[int],
) -> None:
reveal_type(col) # revealed: TypedColumnsClauseRole[int] | SQLCoreOperations[int] | type[int]
```
Next, verify that we can assign `User.id` to a fully specialized version of
`_TypedColumnClauseArgument`:
```py
user_id_as_tcca: _TypedColumnClauseArgument[int] = User.id
```
If we use the generic version of `_TypedColumnClauseArgument` without specialization, we get
`Unknown`:
```py
def extract_t_from_tcca[T](col: _TypedColumnClauseArgument[T]) -> T:
raise NotImplementedError
reveal_type(extract_t_from_tcca(User.id)) # revealed: Unknown
```
However, if we use just the relevant union element of `_TypedColumnClauseArgument`
(`SQLCoreOperations`), it works as expected:
```py
from sqlalchemy.sql.elements import SQLCoreOperations
def extract_t_from_sco[T](col: SQLCoreOperations[T]) -> T:
raise NotImplementedError
reveal_type(extract_t_from_sco(User.id)) # revealed: int
reveal_type(extract_t_from_sco(User.name)) # revealed: str
```
I reported this as <https://github.com/astral-sh/ty/issues/1772>.
Now let's assume we would be able to solve for `T` here. This would mean we would get a type of
`Select[tuple[int, str]]`. Can we use that type and proceed with it? It looks like this works:
```py
from sqlalchemy.sql.selectable import Select
def _(stmt: Select[tuple[int, str]]) -> None:
for row in session.execute(stmt):
reveal_type(row) # revealed: Row[tuple[int, str]]
```
What about the `_tuple` calls? This seems to work:
```py
def _(stmt: Select[tuple[int, str]]) -> None:
result = session.execute(stmt)
reveal_type(result) # revealed: Result[tuple[int, str]]
user = result.one_or_none()
reveal_type(user) # revealed: Row[tuple[int, str]] | None
if not user:
return
reveal_type(user) # revealed: Row[tuple[int, str]] & ~AlwaysFalsy
reveal_type(user._tuple()) # revealed: tuple[int, str]
```
What about `.tuples()`? That seems to work as well:
```py
def _(stmt: Select[tuple[int, str]]) -> None:
for user_id, name in session.execute(stmt).tuples():
reveal_type(user_id) # revealed: int
reveal_type(name) # revealed: str
```
What about the `.scalar` calls? Those seem to work too:
```py
def _(stmt: Select[tuple[int]]) -> None:
user_id = session.scalar(stmt)
reveal_type(user_id) # revealed: int | None
reveal_type(session.scalars(stmt).first()) # revealed: int | None
```

View File

@@ -102,6 +102,38 @@ Other values are invalid.
P4 = ParamSpec("P4", default=int)
```
### `default` parameter in `typing_extensions.ParamSpec`
```toml
[environment]
python-version = "3.12"
```
The `default` parameter to `ParamSpec` is available from `typing_extensions` in Python 3.12 and
earlier.
```py
from typing import ParamSpec
from typing_extensions import ParamSpec as ExtParamSpec
# This shouldn't emit a diagnostic
P1 = ExtParamSpec("P1", default=[int, str])
# But, this should
# error: [invalid-paramspec] "The `default` parameter of `typing.ParamSpec` was added in Python 3.13"
P2 = ParamSpec("P2", default=[int, str])
```
And, it allows the same set of values as `typing.ParamSpec`.
```py
P3 = ExtParamSpec("P3", default=...)
P4 = ExtParamSpec("P4", default=P3)
# error: [invalid-paramspec]
P5 = ExtParamSpec("P5", default=int)
```
### Forward references in stub files
Stubs natively support forward references, so patterns that would raise `NameError` at runtime are

View File

@@ -0,0 +1,19 @@
# `ParamSpec` regression on 3.9
```toml
[environment]
python-version = "3.9"
```
This used to panic when run on Python 3.9 because `ParamSpec` was introduced in Python 3.10 and the
diagnostic message for `invalid-exception-caught` expects to construct `typing.ParamSpec`.
```py
# error: [invalid-syntax]
def foo[**P]() -> None:
try:
pass
# error: [invalid-exception-caught] "Invalid object caught in an exception handler: Object has type `typing.ParamSpec`"
except P:
pass
```

View File

@@ -4168,6 +4168,8 @@ pub enum KnownClass {
SpecialForm,
TypeVar,
ParamSpec,
// typing_extensions.ParamSpec
ExtensionsParamSpec, // must be distinct from typing.ParamSpec, backports new features
ParamSpecArgs,
ParamSpecKwargs,
ProtocolMeta,
@@ -4239,6 +4241,7 @@ impl KnownClass {
| Self::TypeVar
| Self::ExtensionsTypeVar
| Self::ParamSpec
| Self::ExtensionsParamSpec
| Self::ParamSpecArgs
| Self::ParamSpecKwargs
| Self::TypeVarTuple
@@ -4371,6 +4374,7 @@ impl KnownClass {
| KnownClass::TypeVar
| KnownClass::ExtensionsTypeVar
| KnownClass::ParamSpec
| KnownClass::ExtensionsParamSpec
| KnownClass::ParamSpecArgs
| KnownClass::ParamSpecKwargs
| KnownClass::TypeVarTuple
@@ -4457,6 +4461,7 @@ impl KnownClass {
| KnownClass::TypeVar
| KnownClass::ExtensionsTypeVar
| KnownClass::ParamSpec
| KnownClass::ExtensionsParamSpec
| KnownClass::ParamSpecArgs
| KnownClass::ParamSpecKwargs
| KnownClass::TypeVarTuple
@@ -4543,6 +4548,7 @@ impl KnownClass {
| KnownClass::TypeVar
| KnownClass::ExtensionsTypeVar
| KnownClass::ParamSpec
| KnownClass::ExtensionsParamSpec
| KnownClass::ParamSpecArgs
| KnownClass::ParamSpecKwargs
| KnownClass::TypeVarTuple
@@ -4634,6 +4640,7 @@ impl KnownClass {
| Self::TypeVar
| Self::ExtensionsTypeVar
| Self::ParamSpec
| Self::ExtensionsParamSpec
| Self::ParamSpecArgs
| Self::ParamSpecKwargs
| Self::TypeVarTuple
@@ -4733,6 +4740,7 @@ impl KnownClass {
| KnownClass::TypeVar
| KnownClass::ExtensionsTypeVar
| KnownClass::ParamSpec
| KnownClass::ExtensionsParamSpec
| KnownClass::ParamSpecArgs
| KnownClass::ParamSpecKwargs
| KnownClass::ProtocolMeta
@@ -4806,6 +4814,7 @@ impl KnownClass {
Self::TypeVar => "TypeVar",
Self::ExtensionsTypeVar => "TypeVar",
Self::ParamSpec => "ParamSpec",
Self::ExtensionsParamSpec => "ParamSpec",
Self::ParamSpecArgs => "ParamSpecArgs",
Self::ParamSpecKwargs => "ParamSpecKwargs",
Self::TypeVarTuple => "TypeVarTuple",
@@ -5139,11 +5148,18 @@ impl KnownClass {
Self::TypeAliasType
| Self::ExtensionsTypeVar
| Self::TypeVarTuple
| Self::ParamSpec
| Self::ExtensionsParamSpec
| Self::ParamSpecArgs
| Self::ParamSpecKwargs
| Self::Deprecated
| Self::NewType => KnownModule::TypingExtensions,
Self::ParamSpec => {
if Program::get(db).python_version(db) >= PythonVersion::PY310 {
KnownModule::Typing
} else {
KnownModule::TypingExtensions
}
}
Self::NoDefaultType => {
let python_version = Program::get(db).python_version(db);
@@ -5247,6 +5263,7 @@ impl KnownClass {
| Self::TypeVar
| Self::ExtensionsTypeVar
| Self::ParamSpec
| Self::ExtensionsParamSpec
| Self::ParamSpecArgs
| Self::ParamSpecKwargs
| Self::TypeVarTuple
@@ -5337,6 +5354,7 @@ impl KnownClass {
| Self::TypeVar
| Self::ExtensionsTypeVar
| Self::ParamSpec
| Self::ExtensionsParamSpec
| Self::ParamSpecArgs
| Self::ParamSpecKwargs
| Self::TypeVarTuple
@@ -5420,7 +5438,7 @@ impl KnownClass {
"Iterable" => &[Self::Iterable],
"Iterator" => &[Self::Iterator],
"Mapping" => &[Self::Mapping],
"ParamSpec" => &[Self::ParamSpec],
"ParamSpec" => &[Self::ParamSpec, Self::ExtensionsParamSpec],
"ParamSpecArgs" => &[Self::ParamSpecArgs],
"ParamSpecKwargs" => &[Self::ParamSpecKwargs],
"TypeVarTuple" => &[Self::TypeVarTuple],
@@ -5542,6 +5560,8 @@ impl KnownClass {
| Self::TypedDictFallback
| Self::TypeVar
| Self::ExtensionsTypeVar
| Self::ParamSpec
| Self::ExtensionsParamSpec
| Self::NamedTupleLike
| Self::ConstraintSet
| Self::GenericContext
@@ -5555,7 +5575,6 @@ impl KnownClass {
| Self::TypeAliasType
| Self::NoDefaultType
| Self::SupportsIndex
| Self::ParamSpec
| Self::ParamSpecArgs
| Self::ParamSpecKwargs
| Self::TypeVarTuple
@@ -5970,6 +5989,7 @@ mod tests {
KnownClass::Member | KnownClass::Nonmember | KnownClass::StrEnum => {
PythonVersion::PY311
}
KnownClass::ParamSpec => PythonVersion::PY310,
_ => PythonVersion::PY37,
};
(class, version_added)

View File

@@ -5033,9 +5033,15 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
) => {
self.infer_legacy_typevar(target, call_expr, definition, typevar_class)
}
Some(KnownClass::ParamSpec) => {
self.infer_paramspec(target, call_expr, definition)
}
Some(
paramspec_class @ (KnownClass::ParamSpec
| KnownClass::ExtensionsParamSpec),
) => self.infer_legacy_paramspec(
target,
call_expr,
definition,
paramspec_class,
),
Some(KnownClass::NewType) => {
self.infer_newtype_expression(target, call_expr, definition)
}
@@ -5080,11 +5086,12 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
target_ty
}
fn infer_paramspec(
fn infer_legacy_paramspec(
&mut self,
target: &ast::Expr,
call_expr: &ast::ExprCall,
definition: Definition<'db>,
known_class: KnownClass,
) -> Type<'db> {
fn error<'db>(
context: &InferContext<'db, '_>,
@@ -5101,7 +5108,8 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
let db = self.db();
let arguments = &call_expr.arguments;
let assume_all_features = self.in_stub();
let is_typing_extensions = known_class == KnownClass::ExtensionsParamSpec;
let assume_all_features = self.in_stub() || is_typing_extensions;
let python_version = Program::get(db).python_version(db);
let have_features_from =
|version: PythonVersion| assume_all_features || python_version >= version;
@@ -5594,7 +5602,10 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
self.infer_type_expression(&bound.value);
}
if let Some(default) = arguments.find_keyword("default") {
if let Some(KnownClass::ParamSpec) = known_class {
if matches!(
known_class,
Some(KnownClass::ParamSpec | KnownClass::ExtensionsParamSpec)
) {
self.infer_paramspec_default(&default.value);
} else {
self.infer_type_expression(&default.value);
@@ -8440,7 +8451,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
);
}
}
Some(KnownClass::ParamSpec) => {
Some(KnownClass::ParamSpec | KnownClass::ExtensionsParamSpec) => {
if let Some(builder) = self
.context
.report_lint(&INVALID_PARAMSPEC, call_expression)

View File

@@ -349,6 +349,7 @@ pub(super) fn to_lsp_diagnostic(
sub_diagnostic
.annotations()
.iter()
.filter(|annotation| !annotation.is_primary())
.filter_map(|annotation| {
annotation_to_related_information(db, annotation, encoding)
}),