[ty] Refactor TypeInferenceBuilder::infer_subscript_expression_types (#19658)

crates/ty_python_semantic/src/types/class.rs

@@ -4213,14 +4213,14 @@ pub(crate) struct SliceLiteral {
     pub(crate) step: Option<i32>,
 }
 
-impl<'db> Type<'db> {
+impl<'db> ClassType<'db> {
-    /// If this type represents a valid slice literal, returns a [`SliceLiteral`] describing it.
+    /// If this class is a specialization of `slice`, returns a [`SliceLiteral`] describing it.
     /// Otherwise returns `None`.
     ///
-    /// The type must be a specialization of the `slice` builtin type, where the specialized
+    /// The specialization must be one in which the typevars are solved as being statically known
-    /// typevars are statically known integers or `None`.
+    /// integers or `None`.
     pub(crate) fn slice_literal(self, db: &'db dyn Db) -> Option<SliceLiteral> {
-        let ClassType::Generic(alias) = self.into_nominal_instance()?.class else {
+        let ClassType::Generic(alias) = self else {
             return None;
         };
         if !alias.origin(db).is_known(db, KnownClass::Slice) {

crates/ty_python_semantic/src/types/infer.rs

@@ -8365,255 +8365,268 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
         slice_ty: Type<'db>,
         expr_context: ExprContext,
     ) -> Type<'db> {
-        match (value_ty, slice_ty, slice_ty.slice_literal(self.db())) {
+        let db = self.db();
-            (Type::NominalInstance(instance), _, _)
+        let context = &self.context;
-                if instance.class.is_known(self.db(), KnownClass::VersionInfo) =>
+
+        let inferred = match (value_ty, slice_ty) {
+            (Type::NominalInstance(instance), _)
+                if instance.class.is_known(db, KnownClass::VersionInfo) =>
             {
-                self.infer_subscript_expression_types(
+                Some(self.infer_subscript_expression_types(
                     value_node,
-                    Type::version_info_tuple(self.db()),
+                    Type::version_info_tuple(db),
                     slice_ty,
                     expr_context,
-                )
+                ))
             }
 
-            (Type::Union(union), _, _) => union.map(self.db(), |element| {
+            (Type::Union(union), _) => Some(union.map(db, |element| {
                 self.infer_subscript_expression_types(value_node, *element, slice_ty, expr_context)
-            }),
+            })),
 
             // TODO: we can map over the intersection and fold the results back into an intersection,
             // but we need to make sure we avoid emitting a diagnostic if one positive element has a `__getitem__`
             // method but another does not. This means `infer_subscript_expression_types`
             // needs to return a `Result` rather than eagerly emitting diagnostics.
-            (Type::Intersection(_), _, _) => {
+            (Type::Intersection(_), _) => {
-                todo_type!("Subscript expressions on intersections")
+                Some(todo_type!("Subscript expressions on intersections"))
             }
 
             // Ex) Given `("a", "b", "c", "d")[1]`, return `"b"`
-            (Type::Tuple(tuple_ty), Type::IntLiteral(int), _) if i32::try_from(int).is_ok() => {
+            (Type::Tuple(tuple_ty), Type::IntLiteral(i64_int)) => {
-                let tuple = tuple_ty.tuple(self.db());
+                i32::try_from(i64_int).ok().map(|i32_int| {
-                tuple
+                    let tuple = tuple_ty.tuple(db);
-                    .py_index(
+                    tuple.py_index(db, i32_int).unwrap_or_else(|_| {
-                        self.db(),
-                        i32::try_from(int).expect("checked in branch arm"),
-                    )
-                    .unwrap_or_else(|_| {
                         report_index_out_of_bounds(
-                            &self.context,
+                            context,
                             "tuple",
                             value_node.into(),
                             value_ty,
                             tuple.len().display_minimum(),
-                            int,
+                            i64_int,
                         );
                         Type::unknown()
                     })
+                })
             }
 
             // Ex) Given `("a", 1, Null)[0:2]`, return `("a", 1)`
-            (Type::Tuple(tuple_ty), _, Some(SliceLiteral { start, stop, step })) => {
+            (Type::Tuple(tuple_ty), Type::NominalInstance(NominalInstanceType { class, .. })) => {
-                let TupleSpec::Fixed(tuple) = tuple_ty.tuple(self.db()) else {
+                class
+                    .slice_literal(db)
+                    .map(|SliceLiteral { start, stop, step }| {
+                        let TupleSpec::Fixed(tuple) = tuple_ty.tuple(db) else {
                             return todo_type!("slice into variable-length tuple");
                         };
 
-                if let Ok(new_elements) = tuple.py_slice(self.db(), start, stop, step) {
+                        if let Ok(new_elements) = tuple.py_slice(db, start, stop, step) {
-                    TupleType::from_elements(self.db(), new_elements.copied())
+                            TupleType::from_elements(db, new_elements.copied())
                         } else {
-                    report_slice_step_size_zero(&self.context, value_node.into());
+                            report_slice_step_size_zero(context, value_node.into());
                             Type::unknown()
                         }
+                    })
             }
 
             // Ex) Given `"value"[1]`, return `"a"`
-            (Type::StringLiteral(literal_ty), Type::IntLiteral(int), _)
+            (Type::StringLiteral(literal_ty), Type::IntLiteral(i64_int)) => {
-                if i32::try_from(int).is_ok() =>
+                i32::try_from(i64_int).ok().map(|i32_int| {
-            {
+                    let literal_value = literal_ty.value(db);
-                let literal_value = literal_ty.value(self.db());
                     (&mut literal_value.chars())
-                    .py_index(
+                        .py_index(db, i32_int)
-                        self.db(),
+                        .map(|ch| Type::string_literal(db, &ch.to_string()))
-                        i32::try_from(int).expect("checked in branch arm"),
-                    )
-                    .map(|ch| Type::string_literal(self.db(), &ch.to_string()))
                         .unwrap_or_else(|_| {
                             report_index_out_of_bounds(
-                            &self.context,
+                                context,
                                 "string",
                                 value_node.into(),
                                 value_ty,
                                 literal_value.chars().count(),
-                            int,
+                                i64_int,
                             );
                             Type::unknown()
                         })
+                })
             }
 
             // Ex) Given `"value"[1:3]`, return `"al"`
-            (Type::StringLiteral(literal_ty), _, Some(SliceLiteral { start, stop, step })) => {
+            (
-                let literal_value = literal_ty.value(self.db());
+                Type::StringLiteral(literal_ty),
-
+                Type::NominalInstance(NominalInstanceType { class, .. }),
+            ) => class
+                .slice_literal(db)
+                .map(|SliceLiteral { start, stop, step }| {
+                    let literal_value = literal_ty.value(db);
                     let chars: Vec<_> = literal_value.chars().collect();
 
-                if let Ok(new_chars) = chars.py_slice(self.db(), start, stop, step) {
+                    if let Ok(new_chars) = chars.py_slice(db, start, stop, step) {
                         let literal: String = new_chars.collect();
-                    Type::string_literal(self.db(), &literal)
+                        Type::string_literal(db, &literal)
                     } else {
-                    report_slice_step_size_zero(&self.context, value_node.into());
+                        report_slice_step_size_zero(context, value_node.into());
                         Type::unknown()
                     }
-            }
+                }),
 
             // Ex) Given `b"value"[1]`, return `97` (i.e., `ord(b"a")`)
-            (Type::BytesLiteral(literal_ty), Type::IntLiteral(int), _)
+            (Type::BytesLiteral(literal_ty), Type::IntLiteral(i64_int)) => {
-                if i32::try_from(int).is_ok() =>
+                i32::try_from(i64_int).ok().map(|i32_int| {
-            {
+                    let literal_value = literal_ty.value(db);
-                let literal_value = literal_ty.value(self.db());
                     literal_value
-                    .py_index(
+                        .py_index(db, i32_int)
-                        self.db(),
-                        i32::try_from(int).expect("checked in branch arm"),
-                    )
                         .map(|byte| Type::IntLiteral((*byte).into()))
                         .unwrap_or_else(|_| {
                             report_index_out_of_bounds(
-                            &self.context,
+                                context,
                                 "bytes literal",
                                 value_node.into(),
                                 value_ty,
                                 literal_value.len(),
-                            int,
+                                i64_int,
                             );
                             Type::unknown()
                         })
+                })
             }
 
             // Ex) Given `b"value"[1:3]`, return `b"al"`
-            (Type::BytesLiteral(literal_ty), _, Some(SliceLiteral { start, stop, step })) => {
+            (
-                let literal_value = literal_ty.value(self.db());
+                Type::BytesLiteral(literal_ty),
+                Type::NominalInstance(NominalInstanceType { class, .. }),
+            ) => class
+                .slice_literal(db)
+                .map(|SliceLiteral { start, stop, step }| {
+                    let literal_value = literal_ty.value(db);
 
-                if let Ok(new_bytes) = literal_value.py_slice(self.db(), start, stop, step) {
+                    if let Ok(new_bytes) = literal_value.py_slice(db, start, stop, step) {
                         let new_bytes: Vec<u8> = new_bytes.copied().collect();
-                    Type::bytes_literal(self.db(), &new_bytes)
+                        Type::bytes_literal(db, &new_bytes)
                     } else {
-                    report_slice_step_size_zero(&self.context, value_node.into());
+                        report_slice_step_size_zero(context, value_node.into());
                         Type::unknown()
                     }
-            }
+                }),
 
             // Ex) Given `"value"[True]`, return `"a"`
             (
                 Type::Tuple(_) | Type::StringLiteral(_) | Type::BytesLiteral(_),
                 Type::BooleanLiteral(bool),
-                _,
+            ) => Some(self.infer_subscript_expression_types(
-            ) => self.infer_subscript_expression_types(
                 value_node,
                 value_ty,
                 Type::IntLiteral(i64::from(bool)),
                 expr_context,
+            )),
+
+            (Type::SpecialForm(SpecialFormType::Protocol), Type::Tuple(typevars)) => {
+                Some(match typevars.tuple(db) {
+                    TupleSpec::Fixed(typevars) => self
+                        .legacy_generic_class_context(
+                            value_node,
+                            typevars.elements_slice(),
+                            LegacyGenericBase::Protocol,
+                        )
+                        .map(|context| {
+                            Type::KnownInstance(KnownInstanceType::SubscriptedProtocol(context))
+                        })
+                        .unwrap_or_else(Type::unknown),
+                    // TODO: emit a diagnostic
+                    TupleSpec::Variable(_) => Type::unknown(),
+                })
+            }
+
+            (Type::SpecialForm(SpecialFormType::Protocol), typevar) => Some(
+                self.legacy_generic_class_context(
+                    value_node,
+                    std::slice::from_ref(&typevar),
+                    LegacyGenericBase::Protocol,
+                )
+                .map(|context| Type::KnownInstance(KnownInstanceType::SubscriptedProtocol(context)))
+                .unwrap_or_else(Type::unknown),
             ),
 
-            (Type::SpecialForm(SpecialFormType::Protocol), Type::Tuple(typevars), _) => {
+            (Type::KnownInstance(KnownInstanceType::SubscriptedProtocol(_)), _) => {
-                let TupleSpec::Fixed(typevars) = typevars.tuple(self.db()) else {
                 // TODO: emit a diagnostic
-                    return Type::unknown();
+                Some(todo_type!("doubly-specialized typing.Protocol"))
-                };
-                self.legacy_generic_class_context(
-                    value_node,
-                    typevars.elements_slice(),
-                    LegacyGenericBase::Protocol,
-                )
-                .map(|context| Type::KnownInstance(KnownInstanceType::SubscriptedProtocol(context)))
-                .unwrap_or_else(Type::unknown)
             }
 
-            (Type::SpecialForm(SpecialFormType::Protocol), typevar, _) => self
+            (Type::SpecialForm(SpecialFormType::Generic), Type::Tuple(typevars)) => {
+                Some(match typevars.tuple(db) {
+                    TupleSpec::Fixed(typevars) => self
                         .legacy_generic_class_context(
-                    value_node,
-                    std::slice::from_ref(&typevar),
-                    LegacyGenericBase::Protocol,
-                )
-                .map(|context| Type::KnownInstance(KnownInstanceType::SubscriptedProtocol(context)))
-                .unwrap_or_else(Type::unknown),
-
-            (Type::KnownInstance(KnownInstanceType::SubscriptedProtocol(_)), _, _) => {
-                // TODO: emit a diagnostic
-                todo_type!("doubly-specialized typing.Protocol")
-            }
-
-            (Type::SpecialForm(SpecialFormType::Generic), Type::Tuple(typevars), _) => {
-                let TupleSpec::Fixed(typevars) = typevars.tuple(self.db()) else {
-                    // TODO: emit a diagnostic
-                    return Type::unknown();
-                };
-                self.legacy_generic_class_context(
                             value_node,
                             typevars.elements_slice(),
                             LegacyGenericBase::Generic,
                         )
-                .map(|context| Type::KnownInstance(KnownInstanceType::SubscriptedGeneric(context)))
+                        .map(|context| {
-                .unwrap_or_else(Type::unknown)
+                            Type::KnownInstance(KnownInstanceType::SubscriptedGeneric(context))
+                        })
+                        .unwrap_or_else(Type::unknown),
+                    // TODO: emit a diagnostic
+                    TupleSpec::Variable(_) => Type::unknown(),
+                })
             }
 
-            (Type::SpecialForm(SpecialFormType::Generic), typevar, _) => self
+            (Type::SpecialForm(SpecialFormType::Generic), typevar) => Some(
-                .legacy_generic_class_context(
+                self.legacy_generic_class_context(
                     value_node,
                     std::slice::from_ref(&typevar),
                     LegacyGenericBase::Generic,
                 )
                 .map(|context| Type::KnownInstance(KnownInstanceType::SubscriptedGeneric(context)))
                 .unwrap_or_else(Type::unknown),
+            ),
 
-            (Type::KnownInstance(KnownInstanceType::SubscriptedGeneric(_)), _, _) => {
+            (Type::KnownInstance(KnownInstanceType::SubscriptedGeneric(_)), _) => {
                 // TODO: emit a diagnostic
-                todo_type!("doubly-specialized typing.Generic")
+                Some(todo_type!("doubly-specialized typing.Generic"))
             }
 
-            (Type::SpecialForm(special_form), _, _) if special_form.class().is_special_form() => {
+            (Type::SpecialForm(special_form), _) if special_form.class().is_special_form() => {
-                todo_type!("Inference of subscript on special form")
+                Some(todo_type!("Inference of subscript on special form"))
             }
 
-            (Type::KnownInstance(known_instance), _, _)
+            (Type::KnownInstance(known_instance), _)
                 if known_instance.class().is_special_form() =>
             {
-                todo_type!("Inference of subscript on special form")
+                Some(todo_type!("Inference of subscript on special form"))
+            }
+
+            _ => None,
+        };
+
+        if let Some(inferred) = inferred {
+            return inferred;
         }
 
-            (value_ty, slice_ty, _) => {
         // If the class defines `__getitem__`, return its return type.
         //
         // See: https://docs.python.org/3/reference/datamodel.html#class-getitem-versus-getitem
-                match value_ty.try_call_dunder(
+        match value_ty.try_call_dunder(db, "__getitem__", CallArguments::positional([slice_ty])) {
-                    self.db(),
+            Ok(outcome) => return outcome.return_type(db),
-                    "__getitem__",
-                    CallArguments::positional([slice_ty]),
-                ) {
-                    Ok(outcome) => return outcome.return_type(self.db()),
             Err(err @ CallDunderError::PossiblyUnbound { .. }) => {
-                        if let Some(builder) = self
+                if let Some(builder) =
-                            .context
+                    context.report_lint(&POSSIBLY_UNBOUND_IMPLICIT_CALL, value_node)
-                            .report_lint(&POSSIBLY_UNBOUND_IMPLICIT_CALL, value_node)
                 {
                     builder.into_diagnostic(format_args!(
                         "Method `__getitem__` of type `{}` is possibly unbound",
-                                value_ty.display(self.db()),
+                        value_ty.display(db),
                     ));
                 }
 
-                        return err.fallback_return_type(self.db());
+                return err.fallback_return_type(db);
             }
             Err(CallDunderError::CallError(_, bindings)) => {
-                        if let Some(builder) =
+                if let Some(builder) = context.report_lint(&CALL_NON_CALLABLE, value_node) {
-                            self.context.report_lint(&CALL_NON_CALLABLE, value_node)
-                        {
                     builder.into_diagnostic(format_args!(
                         "Method `__getitem__` of type `{}` \
                             is not callable on object of type `{}`",
-                                bindings.callable_type().display(self.db()),
+                        bindings.callable_type().display(db),
-                                value_ty.display(self.db()),
+                        value_ty.display(db),
                     ));
                 }
 
-                        return bindings.return_type(self.db());
+                return bindings.return_type(db);
             }
             Err(CallDunderError::MethodNotAvailable) => {
                 // try `__class_getitem__`
@@ -8629,54 +8642,49 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
         // even if the target version is Python 3.8 or lower,
         // despite the fact that there will be no corresponding `__class_getitem__`
         // method in these `sys.version_info` branches.
-                if value_ty.is_subtype_of(self.db(), KnownClass::Type.to_instance(self.db())) {
+        if value_ty.is_subtype_of(db, KnownClass::Type.to_instance(db)) {
-                    let dunder_class_getitem_method =
+            let dunder_class_getitem_method = value_ty.member(db, "__class_getitem__").place;
-                        value_ty.member(self.db(), "__class_getitem__").place;
 
             match dunder_class_getitem_method {
                 Place::Unbound => {}
                 Place::Type(ty, boundness) => {
                     if boundness == Boundness::PossiblyUnbound {
-                                if let Some(builder) = self
+                        if let Some(builder) =
-                                    .context
+                            context.report_lint(&POSSIBLY_UNBOUND_IMPLICIT_CALL, value_node)
-                                    .report_lint(&POSSIBLY_UNBOUND_IMPLICIT_CALL, value_node)
                         {
                             builder.into_diagnostic(format_args!(
                                 "Method `__class_getitem__` of type `{}` \
                                     is possibly unbound",
-                                        value_ty.display(self.db()),
+                                value_ty.display(db),
                             ));
                         }
                     }
 
-                            match ty.try_call(
+                    match ty.try_call(db, &CallArguments::positional([value_ty, slice_ty])) {
-                                self.db(),
+                        Ok(bindings) => return bindings.return_type(db),
-                                &CallArguments::positional([value_ty, slice_ty]),
-                            ) {
-                                Ok(bindings) => return bindings.return_type(self.db()),
                         Err(CallError(_, bindings)) => {
                             if let Some(builder) =
-                                        self.context.report_lint(&CALL_NON_CALLABLE, value_node)
+                                context.report_lint(&CALL_NON_CALLABLE, value_node)
                             {
                                 builder.into_diagnostic(format_args!(
                                     "Method `__class_getitem__` of type `{}` \
                                         is not callable on object of type `{}`",
-                                            bindings.callable_type().display(self.db()),
+                                    bindings.callable_type().display(db),
-                                            value_ty.display(self.db()),
+                                    value_ty.display(db),
                                 ));
                             }
-                                    return bindings.return_type(self.db());
+                            return bindings.return_type(db);
                         }
                     }
                 }
             }
 
             if let Type::ClassLiteral(class) = value_ty {
-                        if class.is_known(self.db(), KnownClass::Type) {
+                if class.is_known(db, KnownClass::Type) {
-                            return KnownClass::GenericAlias.to_instance(self.db());
+                    return KnownClass::GenericAlias.to_instance(db);
                 }
 
-                        if class.generic_context(self.db()).is_some() {
+                if class.generic_context(db).is_some() {
                     // TODO: specialize the generic class using these explicit type
                     // variable assignments. This branch is only encountered when an
                     // explicit class specialization appears inside of some other subscript
@@ -8689,33 +8697,20 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> {
             }
 
             // TODO: properly handle old-style generics; get rid of this temporary hack
-                    if !value_ty.into_class_literal().is_some_and(|class| {
+            if !value_ty
-                        class
+                .into_class_literal()
-                            .iter_mro(self.db(), None)
+                .is_some_and(|class| class.iter_mro(db, None).contains(&ClassBase::Generic))
-                            .contains(&ClassBase::Generic)
+            {
-                    }) {
+                report_non_subscriptable(context, value_node.into(), value_ty, "__class_getitem__");
-                        report_non_subscriptable(
-                            &self.context,
-                            value_node.into(),
-                            value_ty,
-                            "__class_getitem__",
-                        );
             }
         } else {
             if expr_context != ExprContext::Store {
-                        report_non_subscriptable(
+                report_non_subscriptable(context, value_node.into(), value_ty, "__getitem__");
-                            &self.context,
-                            value_node.into(),
-                            value_ty,
-                            "__getitem__",
-                        );
             }
         }
 
         Type::unknown()
     }
-        }
-    }
 
     fn legacy_generic_class_context(
         &self,