Skip to content

Commit d54f868

Browse files
AlexWaygoodAlexWaygood
committed
[ty] Fix Todo type for starred elements in tuple expressions
1 parent 665f680 commit d54f868

File tree

7 files changed

+154
-33
lines changed

7 files changed

+154
-33
lines changed

‎crates/ty_python_semantic/resources/mdtest/expression/len.md‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -43,13 +43,13 @@ reveal_type(len((1,))) # revealed: Literal[1]
4343
reveal_type(len((1, 2))) # revealed: Literal[2]
4444
reveal_type(len(tuple())) # revealed: Literal[0]
4545

46-
# TODO: Handle star unpacks; Should be: Literal[0]
47-
reveal_type(len((*[],))) # revealed: Literal[1]
46+
# could also be `Literal[0]`, but `int` is accurate
47+
reveal_type(len((*[],))) # revealed: int
4848

4949
# fmt: off
5050

51-
# TODO: Handle star unpacks; Should be: Literal[1]
52-
reveal_type(len( # revealed: Literal[2]
51+
# could also be `Literal[1]`, but `int` is accurate
52+
reveal_type(len( # revealed: int
5353
(
5454
*[],
5555
1,
@@ -58,11 +58,11 @@ reveal_type(len( # revealed: Literal[2]
5858

5959
# fmt: on
6060

61-
# TODO: Handle star unpacks; Should be: Literal[2]
62-
reveal_type(len((*[], 1, 2))) # revealed: Literal[3]
61+
# Could also be `Literal[2]`, but `int` is accurate
62+
reveal_type(len((*[], 1, 2))) # revealed: int
6363

64-
# TODO: Handle star unpacks; Should be: Literal[0]
65-
reveal_type(len((*[], *{}))) # revealed: Literal[2]
64+
# Could also be `Literal[0]`, but `int` is accurate
65+
reveal_type(len((*[], *{}))) # revealed: int
6666
```
6767

6868
Tuple subclasses:

‎crates/ty_python_semantic/resources/mdtest/type_compendium/tuple.md‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -531,4 +531,19 @@ x: list[Literal[1, 2, 3]] = list((1, 2, 3))
531531
reveal_type(x) # revealed: list[Literal[1, 2, 3]]
532532
```
533533

534+
## Tuples with starred elements
535+
536+
```py
537+
x = (1, *range(3), 3)
538+
reveal_type(x) # revealed: tuple[Literal[1], *tuple[int, ...], Literal[3]]
539+
540+
y = 1, 2
541+
542+
reveal_type(("foo", *y)) # revealed: tuple[Literal["foo"], Literal[1], Literal[2]]
543+
544+
aa: tuple[list[int], ...] = ([42], *{[56], [78]}, [100])
545+
reveal_type(aa) # revealed: tuple[list[int], *tuple[list[int], ...], list[int]]
546+
547+
```
548+
534549
[not a singleton type]: https://discuss.python.org/t/should-we-specify-in-the-language-reference-that-the-empty-tuple-is-a-singleton/67957

‎crates/ty_python_semantic/src/types/call/bind.rs‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2941,7 +2941,11 @@ impl<'a, 'db> ArgumentTypeChecker<'a, 'db> {
29412941
) {
29422942
let parameters = self.signature.parameters();
29432943
let parameter = &parameters[parameter_index];
2944-
if let Some(mut expected_ty) = parameter.annotated_type() {
2944+
2945+
// TODO: handle starred annotations, e.g. `*args: *Ts` or `*args: *tuple[int, *tuple[str, ...]]`
2946+
if let Some(mut expected_ty) = parameter.annotated_type()
2947+
&& !parameter.has_starred_annotation()
2948+
{
29452949
if let Some(specialization) = self.specialization {
29462950
argument_type = argument_type.apply_specialization(self.db, specialization);
29472951
expected_ty = expected_ty.apply_specialization(self.db, specialization);

‎crates/ty_python_semantic/src/types/infer.rs‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -377,6 +377,27 @@ impl<'db> TypeContext<'db> {
377377
annotation: self.annotation.map(f),
378378
}
379379
}
380+
381+
pub(crate) fn for_starred_expression(
382+
db: &'db dyn Db,
383+
expected_element_type: Type<'db>,
384+
expr: &ast::ExprStarred,
385+
) -> Self {
386+
match &*expr.value {
387+
ast::Expr::List(_) => Self::new(Some(
388+
KnownClass::List.to_specialized_instance(db, [expected_element_type]),
389+
)),
390+
ast::Expr::Set(_) => Self::new(Some(
391+
KnownClass::Set.to_specialized_instance(db, [expected_element_type]),
392+
)),
393+
ast::Expr::Tuple(_) => {
394+
Self::new(Some(Type::homogeneous_tuple(db, expected_element_type)))
395+
}
396+
// `Iterable[<expected_element_type>]` would work well for an arbitrary other node
397+
// if <https://github.com/astral-sh/ty/issues/1576> is implemented.
398+
_ => Self::default(),
399+
}
400+
}
380401
}
381402

382403
/// Returns the statically-known truthiness of a given expression.

‎crates/ty_python_semantic/src/types/infer/builder.rs‎

Lines changed: 66 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -94,7 +94,9 @@ use crate::types::mro::MroErrorKind;
9494
use crate::types::newtype::NewType;
9595
use crate::types::signatures::Signature;
9696
use crate::types::subclass_of::SubclassOfInner;
97-
use crate::types::tuple::{Tuple, TupleLength, TupleSpec, TupleType};
97+
use crate::types::tuple::{
98+
Tuple, TupleLength, TupleSpec, TupleSpecBuilder, TupleType, VariableLengthTuple,
99+
};
98100
use crate::types::typed_dict::{
99101
TypedDictAssignmentKind, validate_typed_dict_constructor, validate_typed_dict_dict_literal,
100102
validate_typed_dict_key_assignment,
@@ -6926,7 +6928,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
69266928
ast::Expr::If(if_expression) => self.infer_if_expression(if_expression, tcx),
69276929
ast::Expr::Lambda(lambda_expression) => self.infer_lambda_expression(lambda_expression),
69286930
ast::Expr::Call(call_expression) => self.infer_call_expression(call_expression, tcx),
6929-
ast::Expr::Starred(starred) => self.infer_starred_expression(starred),
6931+
ast::Expr::Starred(starred) => self.infer_starred_expression(starred, tcx),
69306932
ast::Expr::Yield(yield_expression) => self.infer_yield_expression(yield_expression),
69316933
ast::Expr::YieldFrom(yield_from) => self.infer_yield_from_expression(yield_from),
69326934
ast::Expr::Await(await_expression) => self.infer_await_expression(await_expression),
@@ -7151,25 +7153,66 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
71517153
)
71527154
});
71537155

7156+
let mut is_homogeneous_tuple_annotation = false;
7157+
71547158
let annotated_tuple = tcx
71557159
.known_specialization(self.db(), KnownClass::Tuple)
71567160
.and_then(|specialization| {
7157-
specialization
7161+
let spec = specialization
71587162
.tuple(self.db())
7159-
.expect("the specialization of `KnownClass::Tuple` must have a tuple spec")
7160-
.resize(self.db(), TupleLength::Fixed(elts.len()))
7161-
.ok()
7163+
.expect("the specialization of `KnownClass::Tuple` must have a tuple spec");
7164+
7165+
if matches!(
7166+
spec,
7167+
Tuple::Variable(VariableLengthTuple { prefix, variable: _, suffix})
7168+
if prefix.is_empty() && suffix.is_empty()
7169+
) {
7170+
is_homogeneous_tuple_annotation = true;
7171+
}
7172+
7173+
spec.resize(self.db(), TupleLength::Fixed(elts.len())).ok()
71627174
});
71637175

71647176
let mut annotated_elt_tys = annotated_tuple.as_ref().map(Tuple::all_elements);
71657177

71667178
let db = self.db();
7167-
let element_types = elts.iter().map(|element| {
7168-
let annotated_elt_ty = annotated_elt_tys.as_mut().and_then(Iterator::next).copied();
7169-
self.infer_expression(element, TypeContext::new(annotated_elt_ty))
7170-
});
71717179

7172-
Type::heterogeneous_tuple(db, element_types)
7180+
let can_use_type_context =
7181+
is_homogeneous_tuple_annotation || elts.iter().all(|elt| !elt.is_starred_expr());
7182+
7183+
let mut infer_element = |elt: &ast::Expr| {
7184+
if can_use_type_context {
7185+
let annotated_elt_ty = annotated_elt_tys.as_mut().and_then(Iterator::next).copied();
7186+
let context = if let ast::Expr::Starred(starred) = elt {
7187+
annotated_elt_ty
7188+
.map(|expected_element_type| {
7189+
TypeContext::for_starred_expression(db, expected_element_type, starred)
7190+
})
7191+
.unwrap_or_default()
7192+
} else {
7193+
TypeContext::new(annotated_elt_ty)
7194+
};
7195+
self.infer_expression(elt, context)
7196+
} else {
7197+
self.infer_expression(elt, TypeContext::default())
7198+
}
7199+
};
7200+
7201+
let mut builder = TupleSpecBuilder::with_capacity(elts.len());
7202+
7203+
for element in elts {
7204+
if element.is_starred_expr() {
7205+
let element_type = infer_element(element);
7206+
// Fine to use `iterate` rather than `try_iterate` here:
7207+
// errors from iterating over something not iterable will have been
7208+
// emitted in the `infer_element` call above.
7209+
builder = builder.concat(db, &element_type.iterate(db));
7210+
} else {
7211+
builder.push(infer_element(element).fallback_to_divergent(db));
7212+
}
7213+
}
7214+
7215+
Type::tuple(TupleType::new(db, &builder.build()))
71737216
}
71747217

71757218
fn infer_list_expression(&mut self, list: &ast::ExprList, tcx: TypeContext<'db>) -> Type<'db> {
@@ -8204,25 +8247,28 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
82048247
}
82058248
}
82068249

8207-
fn infer_starred_expression(&mut self, starred: &ast::ExprStarred) -> Type<'db> {
8250+
fn infer_starred_expression(
8251+
&mut self,
8252+
starred: &ast::ExprStarred,
8253+
tcx: TypeContext<'db>,
8254+
) -> Type<'db> {
82088255
let ast::ExprStarred {
82098256
range: _,
82108257
node_index: _,
82118258
value,
82128259
ctx: _,
82138260
} = starred;
82148261

8215-
let iterable_type = self.infer_expression(value, TypeContext::default());
8262+
let db = self.db();
8263+
let iterable_type = self.infer_expression(value, tcx);
8264+
82168265
iterable_type
8217-
.try_iterate(self.db())
8218-
.map(|tuple| tuple.homogeneous_element_type(self.db()))
8266+
.try_iterate(db)
8267+
.map(|spec| Type::tuple(TupleType::new(db, &spec)))
82198268
.unwrap_or_else(|err| {
82208269
err.report_diagnostic(&self.context, iterable_type, value.as_ref().into());
8221-
err.fallback_element_type(self.db())
8222-
});
8223-
8224-
// TODO
8225-
todo_type!("starred expression")
8270+
Type::homogeneous_tuple(db, err.fallback_element_type(db))
8271+
})
82268272
}
82278273

82288274
fn infer_yield_expression(&mut self, yield_expression: &ast::ExprYield) -> Type<'db> {

‎crates/ty_python_semantic/src/types/infer/builder/annotation_expression.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -122,7 +122,7 @@ impl<'db> TypeInferenceBuilder<'db, '_> {
122122

123123
// Annotation expressions also get special handling for `*args` and `**kwargs`.
124124
ast::Expr::Starred(starred) => {
125-
TypeAndQualifiers::declared(self.infer_starred_expression(starred))
125+
TypeAndQualifiers::declared(self.infer_starred_expression(starred, TypeContext::default()))
126126
}
127127

128128
ast::Expr::BytesLiteral(bytes) => {

‎crates/ty_python_semantic/src/types/signatures.rs‎

Lines changed: 38 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1352,6 +1352,7 @@ impl<'db> Parameters<'db> {
13521352
if let Some(inferred_annotation_type) = inferred_annotation(param) {
13531353
Parameter {
13541354
annotated_type: Some(inferred_annotation_type),
1355+
has_starred_annotation: false,
13551356
inferred_annotation: true,
13561357
kind: ParameterKind::PositionalOnly {
13571358
name: Some(param.parameter.name.id.clone()),
@@ -1396,6 +1397,7 @@ impl<'db> Parameters<'db> {
13961397
if let Some(inferred_annotation_type) = inferred_annotation(arg) {
13971398
Parameter {
13981399
annotated_type: Some(inferred_annotation_type),
1400+
has_starred_annotation: false,
13991401
inferred_annotation: true,
14001402
kind: ParameterKind::PositionalOrKeyword {
14011403
name: arg.parameter.name.id.clone(),
@@ -1591,6 +1593,15 @@ pub(crate) struct Parameter<'db> {
15911593
/// the context, like `Self` for the `self` parameter of instance methods.
15921594
inferred_annotation: bool,
15931595

1596+
/// Variadic parameters can have starred annotations, e.g.
1597+
/// - `*args: *Ts`
1598+
/// - `*args: *tuple[int, ...]`
1599+
/// - `*args: *tuple[int, *tuple[str, ...], bytes]`
1600+
///
1601+
/// The `*` prior to the type gives the annotation a different meaning,
1602+
/// so this must be propagated upwards.
1603+
has_starred_annotation: bool,
1604+
15941605
kind: ParameterKind<'db>,
15951606
pub(crate) form: ParameterForm,
15961607
}
@@ -1599,6 +1610,7 @@ impl<'db> Parameter<'db> {
15991610
pub(crate) fn positional_only(name: Option<Name>) -> Self {
16001611
Self {
16011612
annotated_type: None,
1613+
has_starred_annotation: false,
16021614
inferred_annotation: false,
16031615
kind: ParameterKind::PositionalOnly {
16041616
name,
@@ -1611,6 +1623,7 @@ impl<'db> Parameter<'db> {
16111623
pub(crate) fn positional_or_keyword(name: Name) -> Self {
16121624
Self {
16131625
annotated_type: None,
1626+
has_starred_annotation: false,
16141627
inferred_annotation: false,
16151628
kind: ParameterKind::PositionalOrKeyword {
16161629
name,
@@ -1623,6 +1636,7 @@ impl<'db> Parameter<'db> {
16231636
pub(crate) fn variadic(name: Name) -> Self {
16241637
Self {
16251638
annotated_type: None,
1639+
has_starred_annotation: false,
16261640
inferred_annotation: false,
16271641
kind: ParameterKind::Variadic { name },
16281642
form: ParameterForm::Value,
@@ -1632,6 +1646,7 @@ impl<'db> Parameter<'db> {
16321646
pub(crate) fn keyword_only(name: Name) -> Self {
16331647
Self {
16341648
annotated_type: None,
1649+
has_starred_annotation: false,
16351650
inferred_annotation: false,
16361651
kind: ParameterKind::KeywordOnly {
16371652
name,
@@ -1644,6 +1659,7 @@ impl<'db> Parameter<'db> {
16441659
pub(crate) fn keyword_variadic(name: Name) -> Self {
16451660
Self {
16461661
annotated_type: None,
1662+
has_starred_annotation: false,
16471663
inferred_annotation: false,
16481664
kind: ParameterKind::KeywordVariadic { name },
16491665
form: ParameterForm::Value,
@@ -1683,6 +1699,7 @@ impl<'db> Parameter<'db> {
16831699
annotated_type: self
16841700
.annotated_type
16851701
.map(|ty| ty.apply_type_mapping_impl(db, type_mapping, tcx, visitor)),
1702+
has_starred_annotation: self.has_starred_annotation,
16861703
kind: self
16871704
.kind
16881705
.apply_type_mapping_impl(db, type_mapping, tcx, visitor),
@@ -1702,6 +1719,7 @@ impl<'db> Parameter<'db> {
17021719
) -> Self {
17031720
let Parameter {
17041721
annotated_type,
1722+
has_starred_annotation,
17051723
inferred_annotation,
17061724
kind,
17071725
form,
@@ -1746,6 +1764,7 @@ impl<'db> Parameter<'db> {
17461764

17471765
Self {
17481766
annotated_type: Some(annotated_type),
1767+
has_starred_annotation: *has_starred_annotation,
17491768
inferred_annotation: *inferred_annotation,
17501769
kind,
17511770
form: *form,
@@ -1758,10 +1777,20 @@ impl<'db> Parameter<'db> {
17581777
parameter: &ast::Parameter,
17591778
kind: ParameterKind<'db>,
17601779
) -> Self {
1780+
let annotation = parameter.annotation();
1781+
1782+
let (annotated_type, is_starred) = annotation
1783+
.map(|annotation| {
1784+
(
1785+
Some(definition_expression_type(db, definition, annotation)),
1786+
annotation.is_starred_expr(),
1787+
)
1788+
})
1789+
.unwrap_or((None, false));
1790+
17611791
Self {
1762-
annotated_type: parameter
1763-
.annotation()
1764-
.map(|annotation| definition_expression_type(db, definition, annotation)),
1792+
annotated_type,
1793+
has_starred_annotation: is_starred,
17651794
kind,
17661795
form: ParameterForm::Value,
17671796
inferred_annotation: false,
@@ -1814,6 +1843,12 @@ impl<'db> Parameter<'db> {
18141843
self.annotated_type
18151844
}
18161845

1846+
/// Return `true` if this parameter has a starred annotation,
1847+
/// e.g. `*args: *Ts` or `*args: *tuple[int, *tuple[str, ...], bytes]`
1848+
pub(crate) fn has_starred_annotation(&self) -> bool {
1849+
self.has_starred_annotation
1850+
}
1851+
18171852
/// Kind of the parameter.
18181853
pub(crate) fn kind(&self) -> &ParameterKind<'db> {
18191854
&self.kind

0 commit comments

Comments
 (0)