[ty] add support for mapped union and intersection subscript loads (#18846)

## Summary

Note this modifies the diagnostics a bit. Previously performing
subscript access on something like `NotSubscriptable1 |
NotSubscriptable2` would report the full type as not being
subscriptable:

```
[non-subscriptable] "Cannot subscript object of type `NotSubscriptable1 | NotSubscriptable2` with no `__getitem__` method"
```

Now each erroneous constituent has a separate error:

```
[non-subscriptable] "Cannot subscript object of type `NotSubscriptable2` with no `__getitem__` method"
[non-subscriptable] "Cannot subscript object of type `NotSubscriptable1` with no `__getitem__` method"
```

Closes https://github.com/astral-sh/ty/issues/625

## Test Plan

 mdtest

---------

Co-authored-by: Carl Meyer <carl@astral.sh>
This commit is contained in:
Suneet Tipirneni 2025-06-23 12:38:01 -04:00 committed by GitHub
parent a77db3da3f
commit ef8281b695
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 54 additions and 24 deletions

View file

@ -8125,7 +8125,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
}
fn infer_subscript_expression_types(
&mut self,
&self,
value_node: &ast::Expr,
value_ty: Type<'db>,
slice_ty: Type<'db>,
@ -8140,7 +8140,22 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
slice_ty,
)
}
// If the value type is a union make sure to union the load types.
// For example:
// val: tuple[int] | tuple[str]
// val[0] can be an int or str type
(Type::Union(union_ty), _, _) => union_ty.map(self.db(), |ty| {
self.infer_subscript_expression_types(value_node, *ty, slice_ty)
}),
(Type::Intersection(intersection_ty), _, _) => intersection_ty
.positive(self.db())
.iter()
.map(|ty| self.infer_subscript_expression_types(value_node, *ty, slice_ty))
.fold(
IntersectionBuilder::new(self.db()),
IntersectionBuilder::add_positive,
)
.build(),
// Ex) Given `("a", "b", "c", "d")[1]`, return `"b"`
(Type::Tuple(tuple_ty), Type::IntLiteral(int), _) if i32::try_from(int).is_ok() => {
let tuple = tuple_ty.tuple(self.db());
@ -8446,25 +8461,13 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
);
}
match value_ty {
Type::ClassLiteral(_) => {
// TODO: proper support for generic classes
// For now, just infer `Sequence`, if we see something like `Sequence[str]`. This allows us
// to look up attributes on generic base classes, even if we don't understand generics yet.
// Note that this isn't handled by the clause up above for generic classes
// that use legacy type variables and an explicit `Generic` base class.
// Once we handle legacy typevars, this special case will be removed in
// favor of the specialization logic above.
value_ty
}
_ => Type::unknown(),
}
Type::unknown()
}
}
}
fn legacy_generic_class_context(
&mut self,
&self,
value_node: &ast::Expr,
typevars: &[Type<'db>],
origin: LegacyGenericBase,