[ty] Add precise inference for indexing, slicing and unpacking NamedTuple instances (#19560)

Co-authored-by: Brent Westbrook <[email protected]>

Author: AlexWaygood

crates/ty_python_semantic/resources/mdtest/named_tuple.md

@@ -9,6 +9,7 @@ name, and not just by its numeric position within the tuple:
 
 ```py
 from typing import NamedTuple
+from ty_extensions import static_assert, is_subtype_of, is_assignable_to
 
 class Person(NamedTuple):
     id: int
@@ -24,10 +25,45 @@ reveal_type(alice.id)  # revealed: int
 reveal_type(alice.name)  # revealed: str
 reveal_type(alice.age)  # revealed: int | None
 
-# TODO: These should reveal the types of the fields
-reveal_type(alice[0])  # revealed: Unknown
-reveal_type(alice[1])  # revealed: Unknown
-reveal_type(alice[2])  # revealed: Unknown
+# revealed: tuple[<class 'Person'>, <class 'tuple[int, str, int | None]'>, <class 'Sequence[int | str | None]'>, <class 'Reversible[int | str | None]'>, <class 'Collection[int | str | None]'>, <class 'Iterable[int | str | None]'>, <class 'Container[int | str | None]'>, typing.Protocol, typing.Generic, <class 'object'>]
+reveal_type(Person.__mro__)
+
+static_assert(is_subtype_of(Person, tuple[int, str, int | None]))
+static_assert(is_subtype_of(Person, tuple[object, ...]))
+static_assert(not is_assignable_to(Person, tuple[int, str, int]))
+static_assert(not is_assignable_to(Person, tuple[int, str]))
+
+reveal_type(len(alice))  # revealed: Literal[3]
+reveal_type(bool(alice))  # revealed: Literal[True]
+
+reveal_type(alice[0])  # revealed: int
+reveal_type(alice[1])  # revealed: str
+reveal_type(alice[2])  # revealed: int | None
+
+# error: [index-out-of-bounds] "Index 3 is out of bounds for tuple `Person` with length 3"
+reveal_type(alice[3])  # revealed: Unknown
+
+reveal_type(alice[-1])  # revealed: int | None
+reveal_type(alice[-2])  # revealed: str
+reveal_type(alice[-3])  # revealed: int
+
+# error: [index-out-of-bounds] "Index -4 is out of bounds for tuple `Person` with length 3"
+reveal_type(alice[-4])  # revealed: Unknown
+
+reveal_type(alice[1:])  # revealed: tuple[str, int | None]
+reveal_type(alice[::-1])  # revealed: tuple[int | None, str, int]
+
+alice_id, alice_name, alice_age = alice
+reveal_type(alice_id)  # revealed: int
+reveal_type(alice_name)  # revealed: str
+reveal_type(alice_age)  # revealed: int | None
+
+# error: [invalid-assignment] "Not enough values to unpack: Expected 4"
+a, b, c, d = alice
+# error: [invalid-assignment] "Too many values to unpack: Expected 2"
+a, b = alice
+*_, age = alice
+reveal_type(age)  # revealed: int | None
 
 # error: [missing-argument]
 Person(3)

crates/ty_python_semantic/resources/primer/bad.txt

@@ -1,13 +1,16 @@
 Tanjun  # too many iterations
+altair  # too many iterations (uses packaging)
 antidote  # hangs / slow (single threaded)
 artigraph  # cycle panics (value_type_)
 arviz  # too many iterations on versions of arviz newer than https://github.com/arviz-devs/arviz/commit/3205b82bb4d6097c31f7334d7ac51a6de37002d0
 core  # cycle panics (value_type_)
 cpython  # too many cycle iterations
+graphql-core  # stack overflow
 hydpy  # too many iterations
 ibis  # too many iterations
 jax  # too many iterations
 mypy # too many iterations (self-recursive type alias)
+nox # too many iterations (uses packaging)
 packaging  # too many iterations
 pandas  # slow (9s)
 pandas-stubs  # panics on versions of pandas-stubs newer than https://github.com/pandas-dev/pandas-stubs/commit/bf1221eb7ea0e582c30fe233d1f4f5713fce376b
@@ -21,4 +24,5 @@ setuptools  # vendors packaging, see above
 spack  # slow, success, but mypy-primer hangs processing the output
 spark  # too many iterations
 steam.py  # hangs (single threaded)
+streamlit  # too many iterations (uses packaging)
 xarray  # too many iterations

crates/ty_python_semantic/resources/primer/good.txt

@@ -10,7 +10,6 @@ aioredis
 aiortc
 alectryon
 alerta
-altair
 anyio
 apprise
 async-utils
@@ -41,7 +40,6 @@ flake8
 flake8-pyi
 freqtrade
 git-revise
-graphql-core
 httpx-caching
 hydra-zen
 ignite
@@ -64,7 +62,6 @@ more-itertools
 mypy-protobuf
 mypy_primer
 nionutils
-nox
 openlibrary
 operator
 optuna
@@ -107,7 +104,6 @@ starlette
 static-frame
 stone
 strawberry
-streamlit
 svcs
 sympy
 tornado

crates/ty_python_semantic/src/types.rs

@@ -9389,7 +9389,19 @@ impl<'db> BoundSuperType<'db> {
             ));
         }
 
-        let pivot_class = ClassBase::try_from_type(db, pivot_class_type).ok_or({
+        // TODO: having to get a class-literal just to pass it in here is silly.
+        // `BoundSuperType` should probably not be using `ClassBase::try_from_type` here;
+        // this also leads to false negatives in some cases. See discussion in
+        // <https://github.com/astral-sh/ruff/pull/19560#discussion_r2271570071>.
+        let pivot_class = ClassBase::try_from_type(
+            db,
+            pivot_class_type,
+            KnownClass::Object
+                .to_class_literal(db)
+                .into_class_literal()
+                .expect("`object` should always exist in typeshed"),
+        )
+        .ok_or({
             BoundSuperError::InvalidPivotClassType {
                 pivot_class: pivot_class_type,
             }

crates/ty_python_semantic/src/types/class.rs

@@ -2240,7 +2240,7 @@ impl<'db> ClassLiteral<'db> {
     ///     y: str = "a"
     /// ```
     /// we return a map `{"x": (int, None), "y": (str, Some(Literal["a"]))}`.
-    fn own_fields(
+    pub(super) fn own_fields(
         self,
         db: &'db dyn Db,
         specialization: Option<Specialization<'db>>,

crates/ty_python_semantic/src/types/class_base.rs

@@ -1,8 +1,9 @@
 use crate::Db;
 use crate::types::generics::Specialization;
+use crate::types::tuple::TupleType;
 use crate::types::{
-    ClassType, DynamicType, KnownClass, KnownInstanceType, MroError, MroIterator, SpecialFormType,
-    Type, TypeMapping, TypeTransformer, todo_type,
+    ClassLiteral, ClassType, DynamicType, KnownClass, KnownInstanceType, MroError, MroIterator,
+    SpecialFormType, Type, TypeMapping, TypeTransformer, todo_type,
 };
 
 /// Enumeration of the possible kinds of types we allow in class bases.
@@ -68,15 +69,28 @@ impl<'db> ClassBase<'db> {
     /// Attempt to resolve `ty` into a `ClassBase`.
     ///
     /// Return `None` if `ty` is not an acceptable type for a class base.
-    pub(super) fn try_from_type(db: &'db dyn Db, ty: Type<'db>) -> Option<Self> {
+    pub(super) fn try_from_type(
+        db: &'db dyn Db,
+        ty: Type<'db>,
+        subclass: ClassLiteral<'db>,
+    ) -> Option<Self> {
         match ty {
             Type::Dynamic(dynamic) => Some(Self::Dynamic(dynamic)),
             Type::ClassLiteral(literal) => {
                 if literal.is_known(db, KnownClass::Any) {
                     Some(Self::Dynamic(DynamicType::Any))
                 } else if literal.is_known(db, KnownClass::NamedTuple) {
-                    // TODO: Figure out the tuple spec for the named tuple
-                    Self::try_from_type(db, KnownClass::Tuple.to_class_literal(db))
+                    let fields = subclass.own_fields(db, None);
+                    Self::try_from_type(
+                        db,
+                        TupleType::heterogeneous(
+                            db,
+                            fields.values().map(|field| field.declared_ty),
+                        )?
+                        .to_class_type(db)
+                        .into(),
+                        subclass,
+                    )
                 } else {
                     Some(Self::Class(literal.default_specialization(db)))
                 }
@@ -85,7 +99,7 @@ impl<'db> ClassBase<'db> {
             Type::NominalInstance(instance)
                 if instance.class(db).is_known(db, KnownClass::GenericAlias) =>
             {
-                Self::try_from_type(db, todo_type!("GenericAlias instance"))
+                Self::try_from_type(db, todo_type!("GenericAlias instance"), subclass)
             }
             Type::SubclassOf(subclass_of) => subclass_of
                 .subclass_of()
@@ -95,7 +109,7 @@ impl<'db> ClassBase<'db> {
                 let valid_element = inter
                     .positive(db)
                     .iter()
-                    .find_map(|elem| ClassBase::try_from_type(db, *elem))?;
+                    .find_map(|elem| ClassBase::try_from_type(db, *elem, subclass))?;
 
                 if ty.is_disjoint_from(db, KnownClass::Type.to_instance(db)) {
                     None
@@ -122,7 +136,7 @@ impl<'db> ClassBase<'db> {
                 if union
                     .elements(db)
                     .iter()
-                    .all(|elem| ClassBase::try_from_type(db, *elem).is_some())
+                    .all(|elem| ClassBase::try_from_type(db, *elem, subclass).is_some())
                 {
                     Some(ClassBase::Dynamic(*dynamic))
                 } else {
@@ -135,7 +149,7 @@ impl<'db> ClassBase<'db> {
             // in which case we want to treat `Never` in a forgiving way and silence diagnostics
             Type::Never => Some(ClassBase::unknown()),
 
-            Type::TypeAlias(alias) => Self::try_from_type(db, alias.value_type(db)),
+            Type::TypeAlias(alias) => Self::try_from_type(db, alias.value_type(db), subclass),
 
             Type::PropertyInstance(_)
             | Type::BooleanLiteral(_)
@@ -202,42 +216,44 @@ impl<'db> ClassBase<'db> {
 
                 // TODO: Classes inheriting from `typing.Type` et al. also have `Generic` in their MRO
                 SpecialFormType::Dict => {
-                    Self::try_from_type(db, KnownClass::Dict.to_class_literal(db))
+                    Self::try_from_type(db, KnownClass::Dict.to_class_literal(db), subclass)
                 }
                 SpecialFormType::List => {
-                    Self::try_from_type(db, KnownClass::List.to_class_literal(db))
+                    Self::try_from_type(db, KnownClass::List.to_class_literal(db), subclass)
                 }
                 SpecialFormType::Type => {
-                    Self::try_from_type(db, KnownClass::Type.to_class_literal(db))
+                    Self::try_from_type(db, KnownClass::Type.to_class_literal(db), subclass)
                 }
                 SpecialFormType::Tuple => {
-                    Self::try_from_type(db, KnownClass::Tuple.to_class_literal(db))
+                    Self::try_from_type(db, KnownClass::Tuple.to_class_literal(db), subclass)
                 }
                 SpecialFormType::Set => {
-                    Self::try_from_type(db, KnownClass::Set.to_class_literal(db))
+                    Self::try_from_type(db, KnownClass::Set.to_class_literal(db), subclass)
                 }
                 SpecialFormType::FrozenSet => {
-                    Self::try_from_type(db, KnownClass::FrozenSet.to_class_literal(db))
+                    Self::try_from_type(db, KnownClass::FrozenSet.to_class_literal(db), subclass)
                 }
                 SpecialFormType::ChainMap => {
-                    Self::try_from_type(db, KnownClass::ChainMap.to_class_literal(db))
+                    Self::try_from_type(db, KnownClass::ChainMap.to_class_literal(db), subclass)
                 }
                 SpecialFormType::Counter => {
-                    Self::try_from_type(db, KnownClass::Counter.to_class_literal(db))
+                    Self::try_from_type(db, KnownClass::Counter.to_class_literal(db), subclass)
                 }
                 SpecialFormType::DefaultDict => {
-                    Self::try_from_type(db, KnownClass::DefaultDict.to_class_literal(db))
+                    Self::try_from_type(db, KnownClass::DefaultDict.to_class_literal(db), subclass)
                 }
                 SpecialFormType::Deque => {
-                    Self::try_from_type(db, KnownClass::Deque.to_class_literal(db))
+                    Self::try_from_type(db, KnownClass::Deque.to_class_literal(db), subclass)
                 }
                 SpecialFormType::OrderedDict => {
-                    Self::try_from_type(db, KnownClass::OrderedDict.to_class_literal(db))
+                    Self::try_from_type(db, KnownClass::OrderedDict.to_class_literal(db), subclass)
                 }
                 SpecialFormType::TypedDict => Some(Self::TypedDict),
-                SpecialFormType::Callable => {
-                    Self::try_from_type(db, todo_type!("Support for Callable as a base class"))
-                }
+                SpecialFormType::Callable => Self::try_from_type(
+                    db,
+                    todo_type!("Support for Callable as a base class"),
+                    subclass,
+                ),
             },
         }
     }

crates/ty_python_semantic/src/types/mro.rs

@@ -151,7 +151,7 @@ impl<'db> Mro<'db> {
                             )
                     ) =>
             {
-                ClassBase::try_from_type(db, *single_base).map_or_else(
+                ClassBase::try_from_type(db, *single_base, class.class_literal(db).0).map_or_else(
                     || Err(MroErrorKind::InvalidBases(Box::from([(0, *single_base)]))),
                     |single_base| {
                         if single_base.has_cyclic_mro(db) {
@@ -186,7 +186,7 @@ impl<'db> Mro<'db> {
                             &original_bases[i + 1..],
                         );
                     } else {
-                        match ClassBase::try_from_type(db, *base) {
+                        match ClassBase::try_from_type(db, *base, class.class_literal(db).0) {
                             Some(valid_base) => resolved_bases.push(valid_base),
                             None => invalid_bases.push((i, *base)),
                         }
@@ -253,7 +253,9 @@ impl<'db> Mro<'db> {
                     // `inconsistent-mro` diagnostic (which would be accurate -- but not nearly as
                     // precise!).
                     for (index, base) in original_bases.iter().enumerate() {
-                        let Some(base) = ClassBase::try_from_type(db, *base) else {
+                        let Some(base) =
+                            ClassBase::try_from_type(db, *base, class.class_literal(db).0)
+                        else {
                             continue;
                         };
                         base_to_indices.entry(base).or_default().push(index);