[ty] Fix `Todo` type for starred elements in tuple expressions

This commit is contained in:
Alex Waygood 2025-11-17 12:02:51 +00:00
parent df66946b89
commit 93e6023d9b
8 changed files with 187 additions and 39 deletions

View File

@ -42,6 +42,12 @@ def f[T](x: T, cond: bool) -> T | list[T]:
return x if cond else [x] return x if cond else [x]
l5: int | list[int] = f(1, True) l5: int | list[int] = f(1, True)
a: list[int] = [1, 2, *(3, 4, 5)]
reveal_type(a) # revealed: list[int]
b: list[list[int]] = [[1], [2], *([3], [4])]
reveal_type(b) # revealed: list[list[int]]
``` ```
`typed_dict.py`: `typed_dict.py`:

View File

@ -43,13 +43,13 @@ reveal_type(len((1,))) # revealed: Literal[1]
reveal_type(len((1, 2))) # revealed: Literal[2] reveal_type(len((1, 2))) # revealed: Literal[2]
reveal_type(len(tuple())) # revealed: Literal[0] reveal_type(len(tuple())) # revealed: Literal[0]
# TODO: Handle star unpacks; Should be: Literal[0] # could also be `Literal[0]`, but `int` is accurate
reveal_type(len((*[],))) # revealed: Literal[1] reveal_type(len((*[],))) # revealed: int
# fmt: off # fmt: off
# TODO: Handle star unpacks; Should be: Literal[1] # could also be `Literal[1]`, but `int` is accurate
reveal_type(len( # revealed: Literal[2] reveal_type(len( # revealed: int
( (
*[], *[],
1, 1,
@ -58,11 +58,11 @@ reveal_type(len( # revealed: Literal[2]
# fmt: on # fmt: on
# TODO: Handle star unpacks; Should be: Literal[2] # Could also be `Literal[2]`, but `int` is accurate
reveal_type(len((*[], 1, 2))) # revealed: Literal[3] reveal_type(len((*[], 1, 2))) # revealed: int
# TODO: Handle star unpacks; Should be: Literal[0] # Could also be `Literal[0]`, but `int` is accurate
reveal_type(len((*[], *{}))) # revealed: Literal[2] reveal_type(len((*[], *{}))) # revealed: int
``` ```
Tuple subclasses: Tuple subclasses:

View File

@ -531,4 +531,18 @@ x: list[Literal[1, 2, 3]] = list((1, 2, 3))
reveal_type(x) # revealed: list[Literal[1, 2, 3]] reveal_type(x) # revealed: list[Literal[1, 2, 3]]
``` ```
## Tuples with starred elements
```py
x = (1, *range(3), 3)
reveal_type(x) # revealed: tuple[Literal[1], *tuple[int, ...], Literal[3]]
y = 1, 2
reveal_type(("foo", *y)) # revealed: tuple[Literal["foo"], Literal[1], Literal[2]]
aa: tuple[list[int], ...] = ([42], *{[56], [78]}, [100])
reveal_type(aa) # revealed: tuple[list[int], *tuple[list[int], ...], list[int]]
```
[not a singleton type]: https://discuss.python.org/t/should-we-specify-in-the-language-reference-that-the-empty-tuple-is-a-singleton/67957 [not a singleton type]: https://discuss.python.org/t/should-we-specify-in-the-language-reference-that-the-empty-tuple-is-a-singleton/67957

View File

@ -2974,7 +2974,11 @@ impl<'a, 'db> ArgumentTypeChecker<'a, 'db> {
) { ) {
let parameters = self.signature.parameters(); let parameters = self.signature.parameters();
let parameter = &parameters[parameter_index]; let parameter = &parameters[parameter_index];
if let Some(mut expected_ty) = parameter.annotated_type() {
// TODO: handle starred annotations, e.g. `*args: *Ts` or `*args: *tuple[int, *tuple[str, ...]]`
if let Some(mut expected_ty) = parameter.annotated_type()
&& !parameter.has_starred_annotation()
{
if let Some(specialization) = self.specialization { if let Some(specialization) = self.specialization {
argument_type = argument_type.apply_specialization(self.db, specialization); argument_type = argument_type.apply_specialization(self.db, specialization);
expected_ty = expected_ty.apply_specialization(self.db, specialization); expected_ty = expected_ty.apply_specialization(self.db, specialization);

View File

@ -36,6 +36,7 @@
//! of iterations, so if we fail to converge, Salsa will eventually panic. (This should of course //! of iterations, so if we fail to converge, Salsa will eventually panic. (This should of course
//! be considered a bug.) //! be considered a bug.)
use ruff_python_ast as ast;
use ruff_db::parsed::{ParsedModuleRef, parsed_module}; use ruff_db::parsed::{ParsedModuleRef, parsed_module};
use ruff_text_size::Ranged; use ruff_text_size::Ranged;
use rustc_hash::{FxHashMap, FxHashSet}; use rustc_hash::{FxHashMap, FxHashSet};
@ -384,6 +385,27 @@ impl<'db> TypeContext<'db> {
self.annotation self.annotation
.is_some_and(|ty| ty.is_typealias_special_form()) .is_some_and(|ty| ty.is_typealias_special_form())
} }
pub(crate) fn for_starred_expression(
db: &'db dyn Db,
expected_element_type: Type<'db>,
expr: &ast::ExprStarred,
) -> Self {
match &*expr.value {
ast::Expr::List(_) => Self::new(Some(
KnownClass::List.to_specialized_instance(db, [expected_element_type]),
)),
ast::Expr::Set(_) => Self::new(Some(
KnownClass::Set.to_specialized_instance(db, [expected_element_type]),
)),
ast::Expr::Tuple(_) => {
Self::new(Some(Type::homogeneous_tuple(db, expected_element_type)))
}
// `Iterable[<expected_element_type>]` would work well for an arbitrary other node
// if <https://github.com/astral-sh/ty/issues/1576> is implemented.
_ => Self::default(),
}
}
} }
/// Returns the statically-known truthiness of a given expression. /// Returns the statically-known truthiness of a given expression.

View File

@ -95,7 +95,9 @@ use crate::types::mro::MroErrorKind;
use crate::types::newtype::NewType; use crate::types::newtype::NewType;
use crate::types::signatures::{Parameter, Parameters, Signature}; use crate::types::signatures::{Parameter, Parameters, Signature};
use crate::types::subclass_of::SubclassOfInner; use crate::types::subclass_of::SubclassOfInner;
use crate::types::tuple::{Tuple, TupleLength, TupleSpec, TupleType}; use crate::types::tuple::{
Tuple, TupleLength, TupleSpec, TupleSpecBuilder, TupleType, VariableLengthTuple,
};
use crate::types::typed_dict::{ use crate::types::typed_dict::{
TypedDictAssignmentKind, validate_typed_dict_constructor, validate_typed_dict_dict_literal, TypedDictAssignmentKind, validate_typed_dict_constructor, validate_typed_dict_dict_literal,
validate_typed_dict_key_assignment, validate_typed_dict_key_assignment,
@ -7048,7 +7050,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
ast::Expr::If(if_expression) => self.infer_if_expression(if_expression, tcx), ast::Expr::If(if_expression) => self.infer_if_expression(if_expression, tcx),
ast::Expr::Lambda(lambda_expression) => self.infer_lambda_expression(lambda_expression), ast::Expr::Lambda(lambda_expression) => self.infer_lambda_expression(lambda_expression),
ast::Expr::Call(call_expression) => self.infer_call_expression(call_expression, tcx), ast::Expr::Call(call_expression) => self.infer_call_expression(call_expression, tcx),
ast::Expr::Starred(starred) => self.infer_starred_expression(starred), ast::Expr::Starred(starred) => self.infer_starred_expression(starred, tcx),
ast::Expr::Yield(yield_expression) => self.infer_yield_expression(yield_expression), ast::Expr::Yield(yield_expression) => self.infer_yield_expression(yield_expression),
ast::Expr::YieldFrom(yield_from) => self.infer_yield_from_expression(yield_from), ast::Expr::YieldFrom(yield_from) => self.infer_yield_from_expression(yield_from),
ast::Expr::Await(await_expression) => self.infer_await_expression(await_expression), ast::Expr::Await(await_expression) => self.infer_await_expression(await_expression),
@ -7284,25 +7286,66 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
) )
}); });
let mut is_homogeneous_tuple_annotation = false;
let annotated_tuple = tcx let annotated_tuple = tcx
.known_specialization(self.db(), KnownClass::Tuple) .known_specialization(self.db(), KnownClass::Tuple)
.and_then(|specialization| { .and_then(|specialization| {
specialization let spec = specialization
.tuple(self.db()) .tuple(self.db())
.expect("the specialization of `KnownClass::Tuple` must have a tuple spec") .expect("the specialization of `KnownClass::Tuple` must have a tuple spec");
.resize(self.db(), TupleLength::Fixed(elts.len()))
.ok() if matches!(
spec,
Tuple::Variable(VariableLengthTuple { prefix, variable: _, suffix})
if prefix.is_empty() && suffix.is_empty()
) {
is_homogeneous_tuple_annotation = true;
}
spec.resize(self.db(), TupleLength::Fixed(elts.len())).ok()
}); });
let mut annotated_elt_tys = annotated_tuple.as_ref().map(Tuple::all_elements); let mut annotated_elt_tys = annotated_tuple.as_ref().map(Tuple::all_elements);
let db = self.db(); let db = self.db();
let element_types = elts.iter().map(|element| {
let annotated_elt_ty = annotated_elt_tys.as_mut().and_then(Iterator::next).copied();
self.infer_expression(element, TypeContext::new(annotated_elt_ty))
});
Type::heterogeneous_tuple(db, element_types) let can_use_type_context =
is_homogeneous_tuple_annotation || elts.iter().all(|elt| !elt.is_starred_expr());
let mut infer_element = |elt: &ast::Expr| {
if can_use_type_context {
let annotated_elt_ty = annotated_elt_tys.as_mut().and_then(Iterator::next).copied();
let context = if let ast::Expr::Starred(starred) = elt {
annotated_elt_ty
.map(|expected_element_type| {
TypeContext::for_starred_expression(db, expected_element_type, starred)
})
.unwrap_or_default()
} else {
TypeContext::new(annotated_elt_ty)
};
self.infer_expression(elt, context)
} else {
self.infer_expression(elt, TypeContext::default())
}
};
let mut builder = TupleSpecBuilder::with_capacity(elts.len());
for element in elts {
if element.is_starred_expr() {
let element_type = infer_element(element);
// Fine to use `iterate` rather than `try_iterate` here:
// errors from iterating over something not iterable will have been
// emitted in the `infer_element` call above.
builder = builder.concat(db, &element_type.iterate(db));
} else {
builder.push(infer_element(element));
}
}
Type::tuple(TupleType::new(db, &builder.build()))
} }
fn infer_list_expression(&mut self, list: &ast::ExprList, tcx: TypeContext<'db>) -> Type<'db> { fn infer_list_expression(&mut self, list: &ast::ExprList, tcx: TypeContext<'db>) -> Type<'db> {
@ -7459,7 +7502,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
let inferable = generic_context.inferable_typevars(self.db()); let inferable = generic_context.inferable_typevars(self.db());
// Remove any union elements of that are unrelated to the collection type. // Remove any union elements of the annotation that are unrelated to the collection type.
// //
// For example, we only want the `list[int]` from `annotation: list[int] | None` if // For example, we only want the `list[int]` from `annotation: list[int] | None` if
// `collection_ty` is `list`. // `collection_ty` is `list`.
@ -7499,8 +7542,8 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
} }
let elt_tcxs = match annotated_elt_tys { let elt_tcxs = match annotated_elt_tys {
None => Either::Left(iter::repeat(TypeContext::default())), None => Either::Left(iter::repeat(None)),
Some(tys) => Either::Right(tys.iter().map(|ty| TypeContext::new(Some(*ty)))), Some(tys) => Either::Right(tys.iter().copied().map(Some)),
}; };
for elts in elts { for elts in elts {
@ -7529,6 +7572,14 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
{ {
let Some(elt) = elt else { continue }; let Some(elt) = elt else { continue };
let elt_tcx = if let ast::Expr::Starred(starred) = elt {
elt_tcx
.map(|ty| TypeContext::for_starred_expression(self.db(), ty, starred))
.unwrap_or_default()
} else {
TypeContext::new(elt_tcx)
};
let inferred_elt_ty = infer_elt_expression(self, elt, elt_tcx); let inferred_elt_ty = infer_elt_expression(self, elt, elt_tcx);
// Simplify the inference based on the declared type of the element. // Simplify the inference based on the declared type of the element.
@ -7542,7 +7593,18 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
// unions for large nested list literals, which the constraint solver struggles with. // unions for large nested list literals, which the constraint solver struggles with.
let inferred_elt_ty = inferred_elt_ty.promote_literals(self.db(), elt_tcx); let inferred_elt_ty = inferred_elt_ty.promote_literals(self.db(), elt_tcx);
builder.infer(Type::TypeVar(elt_ty), inferred_elt_ty).ok()?; builder
.infer(
Type::TypeVar(elt_ty),
if elt.is_starred_expr() {
inferred_elt_ty
.iterate(self.db())
.homogeneous_element_type(self.db())
} else {
inferred_elt_ty
},
)
.ok()?;
} }
} }
@ -8359,7 +8421,11 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
} }
} }
fn infer_starred_expression(&mut self, starred: &ast::ExprStarred) -> Type<'db> { fn infer_starred_expression(
&mut self,
starred: &ast::ExprStarred,
tcx: TypeContext<'db>,
) -> Type<'db> {
let ast::ExprStarred { let ast::ExprStarred {
range: _, range: _,
node_index: _, node_index: _,
@ -8367,17 +8433,16 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
ctx: _, ctx: _,
} = starred; } = starred;
let iterable_type = self.infer_expression(value, TypeContext::default()); let db = self.db();
let iterable_type = self.infer_expression(value, tcx);
iterable_type iterable_type
.try_iterate(self.db()) .try_iterate(db)
.map(|tuple| tuple.homogeneous_element_type(self.db())) .map(|spec| Type::tuple(TupleType::new(db, &spec)))
.unwrap_or_else(|err| { .unwrap_or_else(|err| {
err.report_diagnostic(&self.context, iterable_type, value.as_ref().into()); err.report_diagnostic(&self.context, iterable_type, value.as_ref().into());
err.fallback_element_type(self.db()) Type::homogeneous_tuple(db, err.fallback_element_type(db))
}); })
// TODO
todo_type!("starred expression")
} }
fn infer_yield_expression(&mut self, yield_expression: &ast::ExprYield) -> Type<'db> { fn infer_yield_expression(&mut self, yield_expression: &ast::ExprYield) -> Type<'db> {

View File

@ -166,9 +166,9 @@ impl<'db> TypeInferenceBuilder<'db, '_> {
ast::Expr::StringLiteral(string) => self.infer_string_annotation_expression(string), ast::Expr::StringLiteral(string) => self.infer_string_annotation_expression(string),
// Annotation expressions also get special handling for `*args` and `**kwargs`. // Annotation expressions also get special handling for `*args` and `**kwargs`.
ast::Expr::Starred(starred) => { ast::Expr::Starred(starred) => TypeAndQualifiers::declared(
TypeAndQualifiers::declared(self.infer_starred_expression(starred)) self.infer_starred_expression(starred, TypeContext::default()),
} ),
ast::Expr::BytesLiteral(bytes) => { ast::Expr::BytesLiteral(bytes) => {
if let Some(builder) = self if let Some(builder) = self

View File

@ -1499,6 +1499,7 @@ impl<'db> Parameters<'db> {
if let Some(inferred_annotation_type) = inferred_annotation(param) { if let Some(inferred_annotation_type) = inferred_annotation(param) {
Parameter { Parameter {
annotated_type: Some(inferred_annotation_type), annotated_type: Some(inferred_annotation_type),
has_starred_annotation: false,
inferred_annotation: true, inferred_annotation: true,
kind: ParameterKind::PositionalOnly { kind: ParameterKind::PositionalOnly {
name: Some(param.parameter.name.id.clone()), name: Some(param.parameter.name.id.clone()),
@ -1543,6 +1544,7 @@ impl<'db> Parameters<'db> {
if let Some(inferred_annotation_type) = inferred_annotation(arg) { if let Some(inferred_annotation_type) = inferred_annotation(arg) {
Parameter { Parameter {
annotated_type: Some(inferred_annotation_type), annotated_type: Some(inferred_annotation_type),
has_starred_annotation: false,
inferred_annotation: true, inferred_annotation: true,
kind: ParameterKind::PositionalOrKeyword { kind: ParameterKind::PositionalOrKeyword {
name: arg.parameter.name.id.clone(), name: arg.parameter.name.id.clone(),
@ -1738,6 +1740,15 @@ pub(crate) struct Parameter<'db> {
/// the context, like `Self` for the `self` parameter of instance methods. /// the context, like `Self` for the `self` parameter of instance methods.
inferred_annotation: bool, inferred_annotation: bool,
/// Variadic parameters can have starred annotations, e.g.
/// - `*args: *Ts`
/// - `*args: *tuple[int, ...]`
/// - `*args: *tuple[int, *tuple[str, ...], bytes]`
///
/// The `*` prior to the type gives the annotation a different meaning,
/// so this must be propagated upwards.
has_starred_annotation: bool,
kind: ParameterKind<'db>, kind: ParameterKind<'db>,
pub(crate) form: ParameterForm, pub(crate) form: ParameterForm,
} }
@ -1746,6 +1757,7 @@ impl<'db> Parameter<'db> {
pub(crate) fn positional_only(name: Option<Name>) -> Self { pub(crate) fn positional_only(name: Option<Name>) -> Self {
Self { Self {
annotated_type: None, annotated_type: None,
has_starred_annotation: false,
inferred_annotation: false, inferred_annotation: false,
kind: ParameterKind::PositionalOnly { kind: ParameterKind::PositionalOnly {
name, name,
@ -1758,6 +1770,7 @@ impl<'db> Parameter<'db> {
pub(crate) fn positional_or_keyword(name: Name) -> Self { pub(crate) fn positional_or_keyword(name: Name) -> Self {
Self { Self {
annotated_type: None, annotated_type: None,
has_starred_annotation: false,
inferred_annotation: false, inferred_annotation: false,
kind: ParameterKind::PositionalOrKeyword { kind: ParameterKind::PositionalOrKeyword {
name, name,
@ -1770,6 +1783,7 @@ impl<'db> Parameter<'db> {
pub(crate) fn variadic(name: Name) -> Self { pub(crate) fn variadic(name: Name) -> Self {
Self { Self {
annotated_type: None, annotated_type: None,
has_starred_annotation: false,
inferred_annotation: false, inferred_annotation: false,
kind: ParameterKind::Variadic { name }, kind: ParameterKind::Variadic { name },
form: ParameterForm::Value, form: ParameterForm::Value,
@ -1779,6 +1793,7 @@ impl<'db> Parameter<'db> {
pub(crate) fn keyword_only(name: Name) -> Self { pub(crate) fn keyword_only(name: Name) -> Self {
Self { Self {
annotated_type: None, annotated_type: None,
has_starred_annotation: false,
inferred_annotation: false, inferred_annotation: false,
kind: ParameterKind::KeywordOnly { kind: ParameterKind::KeywordOnly {
name, name,
@ -1791,6 +1806,7 @@ impl<'db> Parameter<'db> {
pub(crate) fn keyword_variadic(name: Name) -> Self { pub(crate) fn keyword_variadic(name: Name) -> Self {
Self { Self {
annotated_type: None, annotated_type: None,
has_starred_annotation: false,
inferred_annotation: false, inferred_annotation: false,
kind: ParameterKind::KeywordVariadic { name }, kind: ParameterKind::KeywordVariadic { name },
form: ParameterForm::Value, form: ParameterForm::Value,
@ -1830,6 +1846,7 @@ impl<'db> Parameter<'db> {
annotated_type: self annotated_type: self
.annotated_type .annotated_type
.map(|ty| ty.apply_type_mapping_impl(db, type_mapping, tcx, visitor)), .map(|ty| ty.apply_type_mapping_impl(db, type_mapping, tcx, visitor)),
has_starred_annotation: self.has_starred_annotation,
kind: self kind: self
.kind .kind
.apply_type_mapping_impl(db, type_mapping, tcx, visitor), .apply_type_mapping_impl(db, type_mapping, tcx, visitor),
@ -1849,6 +1866,7 @@ impl<'db> Parameter<'db> {
) -> Self { ) -> Self {
let Parameter { let Parameter {
annotated_type, annotated_type,
has_starred_annotation,
inferred_annotation, inferred_annotation,
kind, kind,
form, form,
@ -1893,6 +1911,7 @@ impl<'db> Parameter<'db> {
Self { Self {
annotated_type: Some(annotated_type), annotated_type: Some(annotated_type),
has_starred_annotation: *has_starred_annotation,
inferred_annotation: *inferred_annotation, inferred_annotation: *inferred_annotation,
kind, kind,
form: *form, form: *form,
@ -1911,6 +1930,7 @@ impl<'db> Parameter<'db> {
inferred_annotation, inferred_annotation,
kind, kind,
form, form,
has_starred_annotation,
} = self; } = self;
let annotated_type = match annotated_type { let annotated_type = match annotated_type {
@ -1975,6 +1995,7 @@ impl<'db> Parameter<'db> {
inferred_annotation: *inferred_annotation, inferred_annotation: *inferred_annotation,
kind, kind,
form: *form, form: *form,
has_starred_annotation: *has_starred_annotation,
}) })
} }
@ -1984,10 +2005,20 @@ impl<'db> Parameter<'db> {
parameter: &ast::Parameter, parameter: &ast::Parameter,
kind: ParameterKind<'db>, kind: ParameterKind<'db>,
) -> Self { ) -> Self {
let annotation = parameter.annotation();
let (annotated_type, is_starred) = annotation
.map(|annotation| {
(
Some(definition_expression_type(db, definition, annotation)),
annotation.is_starred_expr(),
)
})
.unwrap_or((None, false));
Self { Self {
annotated_type: parameter annotated_type,
.annotation() has_starred_annotation: is_starred,
.map(|annotation| definition_expression_type(db, definition, annotation)),
kind, kind,
form: ParameterForm::Value, form: ParameterForm::Value,
inferred_annotation: false, inferred_annotation: false,
@ -2040,6 +2071,12 @@ impl<'db> Parameter<'db> {
self.annotated_type self.annotated_type
} }
/// Return `true` if this parameter has a starred annotation,
/// e.g. `*args: *Ts` or `*args: *tuple[int, *tuple[str, ...], bytes]`
pub(crate) fn has_starred_annotation(&self) -> bool {
self.has_starred_annotation
}
/// Kind of the parameter. /// Kind of the parameter.
pub(crate) fn kind(&self) -> &ParameterKind<'db> { pub(crate) fn kind(&self) -> &ParameterKind<'db> {
&self.kind &self.kind