From 2c7a90d9fd0b6365bd4f72d44e33cdff953664e3 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Sun, 1 Mar 2026 20:24:21 -0500 Subject: [PATCH 1/2] [ty] Move subscript logic out of `builder.rs` --- .../src/types/infer/builder.rs | 1032 +---------------- .../src/types/infer/builder/subscript.rs | 999 ++++++++++++++++ 2 files changed, 1021 insertions(+), 1010 deletions(-) create mode 100644 crates/ty_python_semantic/src/types/infer/builder/subscript.rs diff --git a/crates/ty_python_semantic/src/types/infer/builder.rs b/crates/ty_python_semantic/src/types/infer/builder.rs index ae02b0f0a78da2..cc60203aaf2fd4 100644 --- a/crates/ty_python_semantic/src/types/infer/builder.rs +++ b/crates/ty_python_semantic/src/types/infer/builder.rs @@ -1,8 +1,8 @@ use std::borrow::Cow; -use itertools::{Either, EitherOrBoth, Itertools}; +use itertools::{Either, Itertools}; use ruff_db::diagnostic::{ - Annotation, Diagnostic, DiagnosticId, Severity, Span, SubDiagnostic, SubDiagnosticSeverity, + Annotation, DiagnosticId, Severity, Span, SubDiagnostic, SubDiagnosticSeverity, }; use ruff_db::files::File; use ruff_db::parsed::{ParsedModuleRef, parsed_module}; @@ -60,7 +60,7 @@ use crate::semantic_index::{ place_table, }; use crate::types::builder::RecursivelyDefined; -use crate::types::call::bind::{CallableDescription, MatchingOverloadIndex}; +use crate::types::call::bind::MatchingOverloadIndex; use crate::types::call::{Argument, Binding, Bindings, CallArguments, CallError, CallErrorKind}; use crate::types::class::{ AbstractMethod, ClassLiteral, CodeGeneratorKind, DynamicClassAnchor, DynamicClassLiteral, @@ -82,23 +82,23 @@ use crate::types::diagnostic::{ INVALID_TYPE_FORM, INVALID_TYPE_GUARD_CALL, INVALID_TYPE_GUARD_DEFINITION, INVALID_TYPE_VARIABLE_BOUND, INVALID_TYPE_VARIABLE_CONSTRAINTS, INVALID_TYPE_VARIABLE_DEFAULT, INVALID_TYPED_DICT_HEADER, INVALID_TYPED_DICT_STATEMENT, IncompatibleBases, MISSING_ARGUMENT, - NO_MATCHING_OVERLOAD, NOT_SUBSCRIPTABLE, PARAMETER_ALREADY_ASSIGNED, - POSSIBLY_MISSING_ATTRIBUTE, POSSIBLY_MISSING_IMPLICIT_CALL, POSSIBLY_MISSING_IMPORT, - SUBCLASS_OF_FINAL_CLASS, TOO_MANY_POSITIONAL_ARGUMENTS, TypedDictDeleteErrorKind, - UNDEFINED_REVEAL, UNKNOWN_ARGUMENT, UNRESOLVED_ATTRIBUTE, UNRESOLVED_GLOBAL, UNRESOLVED_IMPORT, - UNRESOLVED_REFERENCE, UNSUPPORTED_DYNAMIC_BASE, UNSUPPORTED_OPERATOR, USELESS_OVERLOAD_BODY, + NO_MATCHING_OVERLOAD, PARAMETER_ALREADY_ASSIGNED, POSSIBLY_MISSING_ATTRIBUTE, + POSSIBLY_MISSING_IMPLICIT_CALL, POSSIBLY_MISSING_IMPORT, SUBCLASS_OF_FINAL_CLASS, + TOO_MANY_POSITIONAL_ARGUMENTS, TypedDictDeleteErrorKind, UNDEFINED_REVEAL, UNKNOWN_ARGUMENT, + UNRESOLVED_ATTRIBUTE, UNRESOLVED_GLOBAL, UNRESOLVED_IMPORT, UNRESOLVED_REFERENCE, + UNSUPPORTED_DYNAMIC_BASE, UNSUPPORTED_OPERATOR, USELESS_OVERLOAD_BODY, hint_if_stdlib_attribute_exists_on_other_versions, hint_if_stdlib_submodule_exists_on_other_versions, report_attempted_protocol_instantiation, report_bad_dunder_set_call, report_bad_frozen_dataclass_inheritance, report_call_to_abstract_method, report_cannot_delete_typed_dict_key, report_cannot_pop_required_field_on_typed_dict, report_conflicting_metaclass_from_bases, report_duplicate_bases, report_implicit_return_type, report_instance_layout_conflict, - report_invalid_arguments_to_annotated, report_invalid_assignment, - report_invalid_attribute_assignment, report_invalid_class_match_pattern, - report_invalid_exception_caught, report_invalid_exception_cause, - report_invalid_exception_raised, report_invalid_exception_tuple_caught, - report_invalid_generator_function_return_type, report_invalid_key_on_typed_dict, - report_invalid_or_unsupported_base, report_invalid_return_type, report_invalid_total_ordering, + report_invalid_assignment, report_invalid_attribute_assignment, + report_invalid_class_match_pattern, report_invalid_exception_caught, + report_invalid_exception_cause, report_invalid_exception_raised, + report_invalid_exception_tuple_caught, report_invalid_generator_function_return_type, + report_invalid_key_on_typed_dict, report_invalid_or_unsupported_base, + report_invalid_return_type, report_invalid_total_ordering, report_invalid_type_checking_constant, report_invalid_type_param_order, report_invalid_typevar_default_reference, report_match_pattern_against_non_runtime_checkable_protocol, @@ -114,16 +114,13 @@ use crate::types::function::{ OverloadLiteral, function_body_kind, is_implicit_classmethod, }; use crate::types::generics::{ - GenericContext, InferableTypeVars, SpecializationBuilder, bind_typevar, - enclosing_generic_contexts, typing_self, + InferableTypeVars, SpecializationBuilder, bind_typevar, enclosing_generic_contexts, typing_self, }; use crate::types::infer::builder::paramspec_validation::validate_paramspec_components; use crate::types::infer::{nearest_enclosing_class, nearest_enclosing_function}; use crate::types::mro::{DynamicMroErrorKind, StaticMroErrorKind}; use crate::types::newtype::NewType; -use crate::types::special_form::AliasSpec; use crate::types::subclass_of::SubclassOfInner; -use crate::types::subscript::{LegacyGenericOrigin, SubscriptError, SubscriptErrorKind}; use crate::types::tuple::{Tuple, TupleLength, TupleSpecBuilder, TupleType}; use crate::types::typed_dict::{ TypedDictAssignmentKind, TypedDictKeyAssignment, validate_typed_dict_constructor, @@ -131,26 +128,27 @@ use crate::types::typed_dict::{ }; use crate::types::visitor::find_over_type; use crate::types::{ - BoundTypeVarIdentity, BoundTypeVarInstance, CallDunderError, CallableBinding, CallableType, - CallableTypeKind, ClassType, DataclassParams, DynamicType, GenericAlias, InternedConstraintSet, - InternedType, IntersectionBuilder, IntersectionType, KnownClass, KnownInstanceType, KnownUnion, + BoundTypeVarIdentity, CallDunderError, CallableBinding, CallableType, CallableTypeKind, + ClassType, DataclassParams, DynamicType, GenericAlias, InternedConstraintSet, InternedType, + IntersectionBuilder, IntersectionType, KnownClass, KnownInstanceType, KnownUnion, LintDiagnosticGuard, LiteralValueTypeKind, ManualPEP695TypeAliasType, MemberLookupPolicy, MetaclassCandidate, PEP695TypeAliasType, ParamSpecAttrKind, Parameter, ParameterForm, Parameters, Signature, SpecialFormType, StaticClassLiteral, SubclassOfType, Truthiness, Type, TypeAliasType, TypeAndQualifiers, TypeContext, TypeQualifiers, TypeVarBoundOrConstraints, TypeVarBoundOrConstraintsEvaluation, TypeVarConstraints, TypeVarDefaultEvaluation, TypeVarIdentity, TypeVarInstance, TypeVarKind, TypeVarVariance, TypedDictType, UnionBuilder, - UnionType, UnionTypeInstance, any_over_type, binding_type, definition_expression_type, - infer_complete_scope_types, infer_scope_types, todo_type, + UnionType, binding_type, definition_expression_type, infer_complete_scope_types, + infer_scope_types, todo_type, }; use crate::types::{CallableTypes, overrides}; use crate::types::{ClassBase, add_inferred_python_version_hint_to_diagnostic}; use crate::unpack::{EvaluationMode, UnpackPosition}; -use crate::{AnalysisSettings, Db, FxIndexSet, FxOrderSet, Program}; +use crate::{AnalysisSettings, Db, FxIndexSet, Program}; mod annotation_expression; mod binary_expressions; mod paramspec_validation; +mod subscript; mod type_expression; use super::comparisons::{self, BinaryComparisonVisitor}; @@ -14565,992 +14563,6 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { ) } - fn infer_subscript_expression(&mut self, subscript: &ast::ExprSubscript) -> Type<'db> { - let ast::ExprSubscript { - value, - slice, - range: _, - node_index: _, - ctx, - } = subscript; - - match ctx { - ExprContext::Load => self.infer_subscript_load(subscript), - ExprContext::Store => { - let value_ty = self.infer_expression(value, TypeContext::default()); - let slice_ty = self.infer_expression(slice, TypeContext::default()); - self.infer_subscript_expression_types(subscript, value_ty, slice_ty, *ctx); - Type::Never - } - ExprContext::Del => { - let value_ty = self.infer_expression(value, TypeContext::default()); - let slice_ty = self.infer_expression(slice, TypeContext::default()); - self.validate_subscript_deletion(subscript, value_ty, slice_ty); - Type::Never - } - ExprContext::Invalid => { - let value_ty = self.infer_expression(value, TypeContext::default()); - let slice_ty = self.infer_expression(slice, TypeContext::default()); - self.infer_subscript_expression_types(subscript, value_ty, slice_ty, *ctx); - Type::unknown() - } - } - } - - fn infer_subscript_load(&mut self, subscript: &ast::ExprSubscript) -> Type<'db> { - let value_ty = self.infer_expression(&subscript.value, TypeContext::default()); - - // If we have an implicit type alias like `MyList = list[T]`, and if `MyList` is being - // used in another implicit type alias like `Numbers = MyList[int]`, then we infer the - // right hand side as a value expression, and need to handle the specialization here. - if value_ty.is_generic_alias() { - return self.infer_explicit_type_alias_specialization(subscript, value_ty, false); - } - - self.infer_subscript_load_impl(value_ty, subscript) - } - - fn infer_subscript_load_impl( - &mut self, - value_ty: Type<'db>, - subscript: &ast::ExprSubscript, - ) -> Type<'db> { - let ast::ExprSubscript { - range: _, - node_index: _, - value: _, - slice, - ctx, - } = subscript; - - let mut constraint_keys = vec![]; - - // If `value` is a valid reference, we attempt type narrowing by assignment. - if !value_ty.is_unknown() { - if let Some(expr) = PlaceExpr::try_from_expr(subscript) { - let (place, keys) = self.infer_place_load( - PlaceExprRef::from(&expr), - ast::ExprRef::Subscript(subscript), - ); - constraint_keys.extend(keys); - if let Place::Defined(DefinedPlace { - ty, - definedness: Definedness::AlwaysDefined, - .. - }) = place.place - { - // Even if we can obtain the subscript type based on the assignments, we still perform default type inference - // (to store the expression type and to report errors). - let slice_ty = self.infer_expression(slice, TypeContext::default()); - self.infer_subscript_expression_types(subscript, value_ty, slice_ty, *ctx); - return ty; - } - } - } - - let tuple_generic_alias = |db: &'db dyn Db, tuple: Option>| { - let tuple = tuple.unwrap_or_else(|| TupleType::homogeneous(db, Type::unknown())); - Type::from(tuple.to_class_type(db)) - }; - - match value_ty { - Type::ClassLiteral(class) => { - // HACK ALERT: If we are subscripting a generic class, short-circuit the rest of the - // subscript inference logic and treat this as an explicit specialization. - // TODO: Move this logic into a custom callable, and update `find_name_in_mro` to return - // this callable as the `__class_getitem__` method on `type`. That probably requires - // updating all of the subscript logic below to use custom callables for all of the _other_ - // special cases, too. - if class.is_tuple(self.db()) { - return tuple_generic_alias( - self.db(), - self.infer_tuple_type_expression(subscript), - ); - } else if class.is_known(self.db(), KnownClass::Type) { - let argument_ty = self.infer_type_expression(slice); - return Type::KnownInstance(KnownInstanceType::TypeGenericAlias( - InternedType::new(self.db(), argument_ty), - )); - } - - if let Some(generic_context) = class.generic_context(self.db()) - && let Some(class) = class.as_static() - { - return self.infer_explicit_class_specialization( - subscript, - value_ty, - class, - generic_context, - ); - } - } - Type::KnownInstance(KnownInstanceType::TypeAliasType(TypeAliasType::ManualPEP695( - _, - ))) => { - let slice_ty = self.infer_expression(slice, TypeContext::default()); - let mut variables = FxOrderSet::default(); - slice_ty.bind_and_find_all_legacy_typevars( - self.db(), - self.typevar_binding_context, - &mut variables, - ); - let generic_context = GenericContext::from_typevar_instances(self.db(), variables); - return Type::Dynamic(DynamicType::UnknownGeneric(generic_context)); - } - Type::KnownInstance(KnownInstanceType::TypeAliasType(type_alias)) => { - if let Some(generic_context) = type_alias.generic_context(self.db()) { - return self.infer_explicit_type_alias_type_specialization( - subscript, - value_ty, - type_alias, - generic_context, - ); - } - } - Type::SpecialForm(special_form) => match special_form { - SpecialFormType::Tuple => { - return tuple_generic_alias( - self.db(), - self.infer_tuple_type_expression(subscript), - ); - } - SpecialFormType::Literal => match self.infer_literal_parameter_type(slice) { - Ok(result) => { - return Type::KnownInstance(KnownInstanceType::Literal(InternedType::new( - self.db(), - result, - ))); - } - Err(nodes) => { - for node in nodes { - let Some(builder) = self.context.report_lint(&INVALID_TYPE_FORM, node) - else { - continue; - }; - builder.into_diagnostic( - "Type arguments for `Literal` must be `None`, \ - a literal value (int, bool, str, or bytes), or an enum member", - ); - } - return Type::unknown(); - } - }, - SpecialFormType::Annotated => { - let ast::Expr::Tuple(ast::ExprTuple { - elts: ref arguments, - .. - }) = **slice - else { - report_invalid_arguments_to_annotated(&self.context, subscript); - - return self.infer_expression(slice, TypeContext::default()); - }; - - if arguments.len() < 2 { - report_invalid_arguments_to_annotated(&self.context, subscript); - } - - let [type_expr, metadata @ ..] = &arguments[..] else { - for argument in arguments { - self.infer_expression(argument, TypeContext::default()); - } - self.store_expression_type(slice, Type::unknown()); - return Type::unknown(); - }; - - for element in metadata { - self.infer_expression(element, TypeContext::default()); - } - - let ty = self.infer_type_expression(type_expr); - - return Type::KnownInstance(KnownInstanceType::Annotated(InternedType::new( - self.db(), - ty, - ))); - } - SpecialFormType::Optional => { - let db = self.db(); - - if matches!(**slice, ast::Expr::Tuple(_)) - && let Some(builder) = - self.context.report_lint(&INVALID_TYPE_FORM, subscript) - { - builder.into_diagnostic(format_args!( - "`typing.Optional` requires exactly one argument" - )); - } - - let ty = self.infer_type_expression(slice); - - // `Optional[None]` is equivalent to `None`: - if ty.is_none(db) { - return ty; - } - - return Type::KnownInstance(KnownInstanceType::UnionType( - UnionTypeInstance::new( - db, - None, - Ok(UnionType::from_two_elements(db, ty, Type::none(db))), - ), - )); - } - SpecialFormType::Union => { - let db = self.db(); - - match **slice { - ast::Expr::Tuple(ref tuple) => { - let mut elements = tuple - .elts - .iter() - .map(|elt| self.infer_type_expression(elt)) - .peekable(); - - let is_empty = elements.peek().is_none(); - let union_type = Type::KnownInstance(KnownInstanceType::UnionType( - UnionTypeInstance::new( - db, - None, - Ok(UnionType::from_elements(db, elements)), - ), - )); - - if is_empty - && let Some(builder) = - self.context.report_lint(&INVALID_TYPE_FORM, subscript) - { - builder.into_diagnostic( - "`typing.Union` requires at least one type argument", - ); - } - - return union_type; - } - _ => { - return self.infer_expression(slice, TypeContext::default()); - } - } - } - SpecialFormType::Type => { - // Similar to the branch above that handles `type[…]`, handle `typing.Type[…]` - let argument_ty = self.infer_type_expression(slice); - return Type::KnownInstance(KnownInstanceType::TypeGenericAlias( - InternedType::new(self.db(), argument_ty), - )); - } - SpecialFormType::Callable => { - let arguments = if let ast::Expr::Tuple(tuple) = &*subscript.slice { - &*tuple.elts - } else { - std::slice::from_ref(&*subscript.slice) - }; - - // TODO: Remove this once we support Concatenate properly. This is necessary - // to avoid a lot of false positives downstream, because we can't represent the typevar- - // specialized `Callable` types yet. - let num_arguments = arguments.len(); - if num_arguments == 2 { - let first_arg = &arguments[0]; - let second_arg = &arguments[1]; - - if first_arg.is_subscript_expr() { - let first_arg_ty = - self.infer_expression(first_arg, TypeContext::default()); - if let Type::Dynamic(DynamicType::UnknownGeneric(generic_context)) = - first_arg_ty - { - let mut variables = generic_context - .variables(self.db()) - .collect::>(); - - let return_ty = - self.infer_expression(second_arg, TypeContext::default()); - return_ty.bind_and_find_all_legacy_typevars( - self.db(), - self.typevar_binding_context, - &mut variables, - ); - - let generic_context = - GenericContext::from_typevar_instances(self.db(), variables); - return Type::Dynamic(DynamicType::UnknownGeneric(generic_context)); - } - - if let Some(builder) = - self.context.report_lint(&INVALID_TYPE_FORM, subscript) - { - builder.into_diagnostic(format_args!( - "The first argument to `Callable` must be either a list of types, \ - ParamSpec, Concatenate, or `...`", - )); - } - return Type::KnownInstance(KnownInstanceType::Callable( - CallableType::unknown(self.db()), - )); - } - } - - let callable = self - .infer_callable_type(subscript) - .as_callable() - .expect("always returns Type::Callable"); - - return Type::KnownInstance(KnownInstanceType::Callable(callable)); - } - SpecialFormType::LegacyStdlibAlias(alias) => { - let AliasSpec { - class, - expected_argument_number, - } = alias.alias_spec(); - - let args = if let ast::Expr::Tuple(t) = &**slice { - &*t.elts - } else { - std::slice::from_ref(&**slice) - }; - - if args.len() != expected_argument_number { - if let Some(builder) = - self.context.report_lint(&INVALID_TYPE_FORM, subscript) - { - let noun = if expected_argument_number == 1 { - "argument" - } else { - "arguments" - }; - builder.into_diagnostic(format_args!( - "`typing.{name}` requires exactly \ - {expected_argument_number} {noun}, got {got}", - name = special_form.name(), - got = args.len() - )); - } - } - - let arg_types: Vec<_> = args - .iter() - .map(|arg| self.infer_type_expression(arg)) - .collect(); - - return class - .to_specialized_class_type(self.db(), arg_types) - .map(Type::from) - .unwrap_or_else(Type::unknown); - } - _ => {} - }, - - Type::KnownInstance( - KnownInstanceType::UnionType(_) - | KnownInstanceType::Annotated(_) - | KnownInstanceType::Callable(_) - | KnownInstanceType::TypeGenericAlias(_), - ) => { - return self.infer_explicit_type_alias_specialization(subscript, value_ty, false); - } - Type::Dynamic(DynamicType::Unknown) => { - let slice_ty = self.infer_expression(slice, TypeContext::default()); - let mut variables = FxOrderSet::default(); - slice_ty.bind_and_find_all_legacy_typevars( - self.db(), - self.typevar_binding_context, - &mut variables, - ); - let generic_context = GenericContext::from_typevar_instances(self.db(), variables); - return Type::Dynamic(DynamicType::UnknownGeneric(generic_context)); - } - _ => {} - } - - let slice_ty = self.infer_expression(slice, TypeContext::default()); - let result_ty = self.infer_subscript_expression_types(subscript, value_ty, slice_ty, *ctx); - self.narrow_expr_with_applicable_constraints(subscript, result_ty, &constraint_keys) - } - - fn infer_explicit_class_specialization( - &mut self, - subscript: &ast::ExprSubscript, - value_ty: Type<'db>, - generic_class: StaticClassLiteral<'db>, - generic_context: GenericContext<'db>, - ) -> Type<'db> { - let db = self.db(); - let specialize = &|types: &[Option>]| { - Type::from(generic_class.apply_specialization(db, |_| { - generic_context.specialize_partial(db, types.iter().copied()) - })) - }; - - self.infer_explicit_callable_specialization( - subscript, - value_ty, - generic_context, - specialize, - ) - } - - fn infer_explicit_type_alias_type_specialization( - &mut self, - subscript: &ast::ExprSubscript, - value_ty: Type<'db>, - generic_type_alias: TypeAliasType<'db>, - generic_context: GenericContext<'db>, - ) -> Type<'db> { - let db = self.db(); - let specialize = &|types: &[Option>]| { - let type_alias = generic_type_alias.apply_specialization(db, |_| { - generic_context.specialize_partial(db, types.iter().copied()) - }); - - Type::KnownInstance(KnownInstanceType::TypeAliasType(type_alias)) - }; - - self.infer_explicit_callable_specialization( - subscript, - value_ty, - generic_context, - specialize, - ) - } - - fn infer_explicit_callable_specialization( - &mut self, - subscript: &ast::ExprSubscript, - value_ty: Type<'db>, - generic_context: GenericContext<'db>, - specialize: &dyn Fn(&[Option>]) -> Type<'db>, - ) -> Type<'db> { - enum ExplicitSpecializationError { - InvalidParamSpec, - UnsatisfiedBound, - UnsatisfiedConstraints, - /// These two errors override the errors above, causing all specializations to be `Unknown`. - MissingTypeVars, - TooManyArguments, - /// This error overrides the errors above, causing the type itself to be `Unknown`. - NonGeneric, - } - - fn add_typevar_definition<'db>( - db: &'db dyn Db, - diagnostic: &mut Diagnostic, - typevar: BoundTypeVarInstance<'db>, - ) { - let Some(definition) = typevar.typevar(db).definition(db) else { - return; - }; - let file = definition.file(db); - let module = parsed_module(db, file).load(db); - let range = definition.focus_range(db, &module).range(); - diagnostic.annotate( - Annotation::secondary(Span::from(file).with_range(range)) - .message("Type variable defined here"), - ); - } - - let db = self.db(); - let constraints = ConstraintSetBuilder::new(); - let slice_node = subscript.slice.as_ref(); - - let exactly_one_paramspec = generic_context.exactly_one_paramspec(db); - let (type_arguments, store_inferred_type_arguments) = match slice_node { - ast::Expr::Tuple(tuple) => { - if exactly_one_paramspec && !tuple.elts.is_empty() { - (std::slice::from_ref(slice_node), false) - } else { - (tuple.elts.as_slice(), true) - } - } - _ => (std::slice::from_ref(slice_node), false), - }; - let mut inferred_type_arguments = Vec::with_capacity(type_arguments.len()); - - let typevars = generic_context.variables(db); - let typevars_len = typevars.len(); - - let mut specialization_types = Vec::with_capacity(typevars_len); - let mut typevar_with_defaults = 0; - let mut missing_typevars = vec![]; - let mut first_excess_type_argument_index = None; - - // Helper to get the AST node corresponding to the type argument at `index`. - let get_node = |index: usize| -> ast::AnyNodeRef<'_> { - match slice_node { - ast::Expr::Tuple(ast::ExprTuple { elts, .. }) if !exactly_one_paramspec => elts - .get(index) - .expect("type argument index should not be out of range") - .into(), - _ => slice_node.into(), - } - }; - - let mut error: Option = None; - - for (index, item) in typevars.zip_longest(type_arguments.iter()).enumerate() { - match item { - EitherOrBoth::Both(typevar, expr) => { - if typevar.default_type(db).is_some() { - typevar_with_defaults += 1; - } - - let provided_type = if typevar.is_paramspec(db) { - match self.infer_paramspec_explicit_specialization_value( - expr, - exactly_one_paramspec, - ) { - Ok(paramspec_value) => paramspec_value, - Err(()) => { - error = Some(ExplicitSpecializationError::InvalidParamSpec); - Type::paramspec_value_callable(db, Parameters::unknown()) - } - } - } else { - self.infer_type_expression(expr) - }; - - inferred_type_arguments.push(provided_type); - - // TODO consider just accepting the given specialization without checking - // against bounds/constraints, but recording the expression for deferred - // checking at end of scope. This would avoid a lot of cycles caused by eagerly - // doing assignment checks here. - match typevar.typevar(db).bound_or_constraints(db) { - Some(TypeVarBoundOrConstraints::UpperBound(bound)) => { - if provided_type - .when_assignable_to( - db, - bound, - &constraints, - InferableTypeVars::None, - ) - .is_never_satisfied(db) - { - let node = get_node(index); - if let Some(builder) = - self.context.report_lint(&INVALID_TYPE_ARGUMENTS, node) - { - let mut diagnostic = builder.into_diagnostic(format_args!( - "Type `{}` is not assignable to upper bound `{}` \ - of type variable `{}`", - provided_type.display(db), - bound.display(db), - typevar.identity(db).display(db), - )); - add_typevar_definition(db, &mut diagnostic, typevar); - } - error = Some(ExplicitSpecializationError::UnsatisfiedBound); - specialization_types.push(Some(Type::unknown())); - } else { - specialization_types.push(Some(provided_type)); - } - } - Some(TypeVarBoundOrConstraints::Constraints(typevar_constraints)) => { - // TODO: this is wrong, the given specialization needs to be assignable - // to _at least one_ of the individual constraints, not to the union of - // all of them. `int | str` is not a valid specialization of a typevar - // constrained to `(int, str)`. - if provided_type - .when_assignable_to( - db, - typevar_constraints.as_type(db), - &constraints, - InferableTypeVars::None, - ) - .is_never_satisfied(db) - { - let node = get_node(index); - if let Some(builder) = - self.context.report_lint(&INVALID_TYPE_ARGUMENTS, node) - { - let mut diagnostic = builder.into_diagnostic(format_args!( - "Type `{}` does not satisfy constraints `{}` \ - of type variable `{}`", - provided_type.display(db), - typevar_constraints - .elements(db) - .iter() - .map(|c| c.display(db)) - .format("`, `"), - typevar.identity(db).display(db), - )); - add_typevar_definition(db, &mut diagnostic, typevar); - } - error = Some(ExplicitSpecializationError::UnsatisfiedConstraints); - specialization_types.push(Some(Type::unknown())); - } else { - specialization_types.push(Some(provided_type)); - } - } - None => { - specialization_types.push(Some(provided_type)); - } - } - } - EitherOrBoth::Left(typevar) => { - if typevar.default_type(db).is_none() { - // This is an error case, so no need to push into the specialization types. - missing_typevars.push(typevar); - } else { - typevar_with_defaults += 1; - specialization_types.push(None); - } - } - EitherOrBoth::Right(expr) => { - inferred_type_arguments.push(self.infer_type_expression(expr)); - first_excess_type_argument_index.get_or_insert(index); - } - } - } - - if !missing_typevars.is_empty() { - if let Some(builder) = self.context.report_lint(&INVALID_TYPE_ARGUMENTS, subscript) { - let description = CallableDescription::new(db, value_ty); - let s = if missing_typevars.len() > 1 { "s" } else { "" }; - builder.into_diagnostic(format_args!( - "No type argument{s} provided for required type variable{s} `{}`{}", - missing_typevars - .iter() - .map(|tv| tv.typevar(db).name(db)) - .format("`, `"), - if let Some(CallableDescription { kind, name }) = description { - format!(" of {kind} `{name}`") - } else { - String::new() - } - )); - } - error = Some(ExplicitSpecializationError::MissingTypeVars); - } - - if let Some(first_excess_type_argument_index) = first_excess_type_argument_index { - if let Type::GenericAlias(alias) = value_ty - && let spec = alias.specialization(self.db()) - && spec - .types(self.db()) - .contains(&Type::Dynamic(DynamicType::TodoTypeVarTuple)) - { - // Avoid false-positive errors when specializing a class - // that's generic over a legacy TypeVarTuple - } else if typevars_len == 0 { - // Type parameter list cannot be empty, so if we reach here, `value_ty` is not a generic type. - if let Some(builder) = self.context.report_lint(&NOT_SUBSCRIPTABLE, subscript) { - let mut diagnostic = builder.into_diagnostic(format_args!( - "Cannot subscript non-generic type `{}`", - value_ty.display(db) - )); - if match value_ty { - Type::GenericAlias(_) => true, - Type::KnownInstance(KnownInstanceType::UnionType(union)) => union - .value_expression_types(db) - .is_ok_and(|mut tys| tys.any(|ty| ty.is_generic_alias())), - _ => false, - } { - diagnostic.annotate( - self.context - .secondary(&*subscript.value) - .message("Type is already specialized"), - ); - } - } - error = Some(ExplicitSpecializationError::NonGeneric); - } else { - let node = get_node(first_excess_type_argument_index); - if let Some(builder) = self.context.report_lint(&INVALID_TYPE_ARGUMENTS, node) { - let description = CallableDescription::new(db, value_ty); - builder.into_diagnostic(format_args!( - "Too many type arguments{}: expected {}, got {}", - if let Some(CallableDescription { kind, name }) = description { - format!(" to {kind} `{name}`") - } else { - String::new() - }, - if typevar_with_defaults == 0 { - format!("{typevars_len}") - } else { - format!( - "between {} and {}", - typevars_len - typevar_with_defaults, - typevars_len - ) - }, - type_arguments.len(), - )); - } - error = Some(ExplicitSpecializationError::TooManyArguments); - } - } - - if store_inferred_type_arguments { - self.store_expression_type( - slice_node, - Type::heterogeneous_tuple(db, inferred_type_arguments), - ); - } - - match error { - Some(ExplicitSpecializationError::NonGeneric) => Type::unknown(), - Some( - ExplicitSpecializationError::MissingTypeVars - | ExplicitSpecializationError::TooManyArguments, - ) => { - let unknowns = generic_context - .variables(self.db()) - .map(|typevar| { - Some(if typevar.is_paramspec(db) { - Type::paramspec_value_callable(db, Parameters::unknown()) - } else { - Type::unknown() - }) - }) - .collect::>(); - specialize(&unknowns) - } - Some( - ExplicitSpecializationError::UnsatisfiedBound - | ExplicitSpecializationError::UnsatisfiedConstraints - | ExplicitSpecializationError::InvalidParamSpec, - ) - | None => specialize(&specialization_types), - } - } - - fn infer_subscript_expression_types( - &self, - subscript: &ast::ExprSubscript, - value_ty: Type<'db>, - slice_ty: Type<'db>, - expr_context: ExprContext, - ) -> Type<'db> { - #[derive(Debug, Clone, Copy, PartialEq, Eq)] - enum LegacyGenericContextError<'db> { - /// It's invalid to subscript `Generic` or `Protocol` with this type. - InvalidArgument(Type<'db>), - /// It's invalid to subscript `Generic` or `Protocol` with a variadic tuple type. - /// We should emit a diagnostic for this, but we don't yet. - VariadicTupleArguments, - /// It's valid to subscribe `Generic` or `Protocol` with this type, - /// but the type is not yet supported. - NotYetSupported, - /// A duplicate typevar was provided. - DuplicateTypevar(&'db str), - /// A `TypeVarTuple` was provided but not unpacked. - TypeVarTupleMustBeUnpacked, - } - - impl<'db> LegacyGenericContextError<'db> { - const fn into_type(self) -> Type<'db> { - match self { - LegacyGenericContextError::InvalidArgument(_) - | LegacyGenericContextError::VariadicTupleArguments - | LegacyGenericContextError::DuplicateTypevar(_) - | LegacyGenericContextError::TypeVarTupleMustBeUnpacked => Type::unknown(), - LegacyGenericContextError::NotYetSupported => { - todo_type!("ParamSpecs and TypeVarTuples") - } - } - } - } - - let db = self.db(); - - let legacy_generic_class_context = - |typevars: Type<'db>| -> Result, LegacyGenericContextError<'db>> { - let typevars_class_tuple_spec = typevars.exact_tuple_instance_spec(db); - - let typevars = if let Some(tuple_spec) = typevars_class_tuple_spec.as_deref() { - match tuple_spec { - Tuple::Fixed(typevars) => typevars.elements_slice(), - Tuple::Variable(_) => { - return Err(LegacyGenericContextError::VariadicTupleArguments); - } - } - } else { - std::slice::from_ref(&typevars) - }; - - let mut validated_typevars = FxOrderSet::default(); - for ty in typevars { - let argument_ty = *ty; - if let Type::KnownInstance(KnownInstanceType::TypeVar(typevar)) = argument_ty { - let bound = bind_typevar( - db, - self.index, - self.scope().file_scope_id(db), - self.typevar_binding_context, - typevar, - ) - .ok_or(LegacyGenericContextError::InvalidArgument(argument_ty))?; - if !validated_typevars.insert(bound) { - return Err(LegacyGenericContextError::DuplicateTypevar( - typevar.name(db), - )); - } - } else if let Type::NominalInstance(instance) = argument_ty - && instance.has_known_class(db, KnownClass::TypeVarTuple) - { - return Err(LegacyGenericContextError::TypeVarTupleMustBeUnpacked); - } else if any_over_type(db, argument_ty, true, |inner_ty| match inner_ty { - Type::Dynamic( - DynamicType::TodoUnpack | DynamicType::TodoStarredExpression, - ) => true, - Type::NominalInstance(nominal) => { - nominal.has_known_class(db, KnownClass::TypeVarTuple) - } - _ => false, - }) { - return Err(LegacyGenericContextError::NotYetSupported); - } else { - return Err(LegacyGenericContextError::InvalidArgument(argument_ty)); - } - } - Ok(GenericContext::from_typevar_instances( - db, - validated_typevars, - )) - }; - - // Special typing forms for which subscriptions are context-dependent are parsed here, - // outside of `Type::subscript`, which is a pure function that doesn't depend on the - // semantic index or any context-dependent state. - let subscript_result = match value_ty { - Type::SpecialForm(SpecialFormType::Generic) => { - match legacy_generic_class_context(slice_ty) { - Ok(context) => Ok(Type::KnownInstance(KnownInstanceType::SubscriptedGeneric( - context, - ))), - Err(LegacyGenericContextError::InvalidArgument(argument_ty)) => { - Err(SubscriptError::new( - Type::unknown(), - SubscriptErrorKind::InvalidLegacyGenericArgument { - origin: LegacyGenericOrigin::Generic, - argument_ty, - }, - )) - } - Err(LegacyGenericContextError::DuplicateTypevar(typevar_name)) => { - Err(SubscriptError::new( - Type::unknown(), - SubscriptErrorKind::DuplicateTypevar { - origin: LegacyGenericOrigin::Generic, - typevar_name, - }, - )) - } - Err(LegacyGenericContextError::TypeVarTupleMustBeUnpacked) => { - Err(SubscriptError::new( - Type::unknown(), - SubscriptErrorKind::TypeVarTupleNotUnpacked { - origin: LegacyGenericOrigin::Generic, - }, - )) - } - Err( - error @ (LegacyGenericContextError::NotYetSupported - | LegacyGenericContextError::VariadicTupleArguments), - ) => Ok(error.into_type()), - } - } - Type::SpecialForm(SpecialFormType::Protocol) => { - match legacy_generic_class_context(slice_ty) { - Ok(context) => Ok(Type::KnownInstance(KnownInstanceType::SubscriptedProtocol( - context, - ))), - Err(LegacyGenericContextError::InvalidArgument(argument_ty)) => { - Err(SubscriptError::new( - Type::unknown(), - SubscriptErrorKind::InvalidLegacyGenericArgument { - origin: LegacyGenericOrigin::Protocol, - argument_ty, - }, - )) - } - Err(LegacyGenericContextError::DuplicateTypevar(typevar_name)) => { - Err(SubscriptError::new( - Type::unknown(), - SubscriptErrorKind::DuplicateTypevar { - origin: LegacyGenericOrigin::Protocol, - typevar_name, - }, - )) - } - Err(LegacyGenericContextError::TypeVarTupleMustBeUnpacked) => { - Err(SubscriptError::new( - Type::unknown(), - SubscriptErrorKind::TypeVarTupleNotUnpacked { - origin: LegacyGenericOrigin::Protocol, - }, - )) - } - Err( - error @ (LegacyGenericContextError::NotYetSupported - | LegacyGenericContextError::VariadicTupleArguments), - ) => Ok(error.into_type()), - } - } - Type::SpecialForm(SpecialFormType::Concatenate) => { - // TODO: Add proper support for `Concatenate` - let mut variables = FxOrderSet::default(); - slice_ty.bind_and_find_all_legacy_typevars( - db, - self.typevar_binding_context, - &mut variables, - ); - let generic_context = GenericContext::from_typevar_instances(db, variables); - Ok(Type::Dynamic(DynamicType::UnknownGeneric(generic_context))) - } - _ => value_ty.subscript(self.db(), slice_ty, expr_context), - }; - - subscript_result.unwrap_or_else(|e| { - e.report_diagnostics(&self.context, subscript); - e.result_type() - }) - } - - fn infer_slice_expression(&mut self, slice: &ast::ExprSlice) -> Type<'db> { - enum SliceArg<'db> { - Arg(Type<'db>), - Unsupported, - } - - let ast::ExprSlice { - range: _, - node_index: _, - lower, - upper, - step, - } = slice; - - let ty_lower = self.infer_optional_expression(lower.as_deref(), TypeContext::default()); - let ty_upper = self.infer_optional_expression(upper.as_deref(), TypeContext::default()); - let ty_step = self.infer_optional_expression(step.as_deref(), TypeContext::default()); - - let type_to_slice_argument = |ty: Option>| match ty { - Some(ty @ Type::LiteralValue(literal)) if literal.is_int() || literal.is_bool() => { - SliceArg::Arg(ty) - } - Some(ty @ Type::NominalInstance(instance)) - if instance.has_known_class(self.db(), KnownClass::NoneType) => - { - SliceArg::Arg(ty) - } - None => SliceArg::Arg(Type::none(self.db())), - _ => SliceArg::Unsupported, - }; - - match ( - type_to_slice_argument(ty_lower), - type_to_slice_argument(ty_upper), - type_to_slice_argument(ty_step), - ) { - (SliceArg::Arg(lower), SliceArg::Arg(upper), SliceArg::Arg(step)) => { - KnownClass::Slice.to_specialized_instance(self.db(), &[lower, upper, step]) - } - _ => KnownClass::Slice.to_instance(self.db()), - } - } - fn infer_type_parameters(&mut self, type_parameters: &ast::TypeParams) { let ast::TypeParams { range: _, diff --git a/crates/ty_python_semantic/src/types/infer/builder/subscript.rs b/crates/ty_python_semantic/src/types/infer/builder/subscript.rs new file mode 100644 index 00000000000000..474b668d2e0bd6 --- /dev/null +++ b/crates/ty_python_semantic/src/types/infer/builder/subscript.rs @@ -0,0 +1,999 @@ +use itertools::{EitherOrBoth, Itertools}; +use ruff_db::diagnostic::{Annotation, Diagnostic, Span}; +use ruff_db::parsed::parsed_module; +use ruff_python_ast::{self as ast, ExprContext}; +use ruff_text_size::Ranged; + +use super::TypeInferenceBuilder; +use crate::place::{DefinedPlace, Definedness, Place}; +use crate::semantic_index::SemanticIndex; +use crate::semantic_index::definition::Definition; +use crate::semantic_index::place::{PlaceExpr, PlaceExprRef}; +use crate::semantic_index::scope::FileScopeId; +use crate::types::call::bind::CallableDescription; +use crate::types::constraints::ConstraintSetBuilder; +use crate::types::diagnostic::{ + INVALID_TYPE_ARGUMENTS, INVALID_TYPE_FORM, NOT_SUBSCRIPTABLE, + report_invalid_arguments_to_annotated, +}; +use crate::types::generics::{GenericContext, InferableTypeVars, bind_typevar}; +use crate::types::special_form::AliasSpec; +use crate::types::subscript::{LegacyGenericOrigin, SubscriptError, SubscriptErrorKind}; +use crate::types::tuple::{Tuple, TupleType}; +use crate::types::{ + BoundTypeVarInstance, CallableType, DynamicType, InternedType, KnownClass, KnownInstanceType, + Parameters, SpecialFormType, StaticClassLiteral, Type, TypeAliasType, TypeContext, + TypeVarBoundOrConstraints, UnionType, UnionTypeInstance, any_over_type, todo_type, +}; +use crate::{Db, FxOrderSet}; + +impl<'db> TypeInferenceBuilder<'db, '_> { + pub(super) fn infer_subscript_expression( + &mut self, + subscript: &ast::ExprSubscript, + ) -> Type<'db> { + let ast::ExprSubscript { + value, + slice, + range: _, + node_index: _, + ctx, + } = subscript; + + match ctx { + ExprContext::Load => self.infer_subscript_load(subscript), + ExprContext::Store => { + let value_ty = self.infer_expression(value, TypeContext::default()); + let slice_ty = self.infer_expression(slice, TypeContext::default()); + self.infer_subscript_expression_types(subscript, value_ty, slice_ty, *ctx); + Type::Never + } + ExprContext::Del => { + let value_ty = self.infer_expression(value, TypeContext::default()); + let slice_ty = self.infer_expression(slice, TypeContext::default()); + self.validate_subscript_deletion(subscript, value_ty, slice_ty); + Type::Never + } + ExprContext::Invalid => { + let value_ty = self.infer_expression(value, TypeContext::default()); + let slice_ty = self.infer_expression(slice, TypeContext::default()); + self.infer_subscript_expression_types(subscript, value_ty, slice_ty, *ctx); + Type::unknown() + } + } + } + + pub(super) fn infer_subscript_load(&mut self, subscript: &ast::ExprSubscript) -> Type<'db> { + let value_ty = self.infer_expression(&subscript.value, TypeContext::default()); + + // If we have an implicit type alias like `MyList = list[T]`, and if `MyList` is being + // used in another implicit type alias like `Numbers = MyList[int]`, then we infer the + // right hand side as a value expression, and need to handle the specialization here. + if value_ty.is_generic_alias() { + return self.infer_explicit_type_alias_specialization(subscript, value_ty, false); + } + + self.infer_subscript_load_impl(value_ty, subscript) + } + + pub(super) fn infer_subscript_load_impl( + &mut self, + value_ty: Type<'db>, + subscript: &ast::ExprSubscript, + ) -> Type<'db> { + let ast::ExprSubscript { + range: _, + node_index: _, + value: _, + slice, + ctx, + } = subscript; + + let mut constraint_keys = vec![]; + + // If `value` is a valid reference, we attempt type narrowing by assignment. + if !value_ty.is_unknown() { + if let Some(expr) = PlaceExpr::try_from_expr(subscript) { + let (place, keys) = self.infer_place_load( + PlaceExprRef::from(&expr), + ast::ExprRef::Subscript(subscript), + ); + constraint_keys.extend(keys); + if let Place::Defined(DefinedPlace { + ty, + definedness: Definedness::AlwaysDefined, + .. + }) = place.place + { + // Even if we can obtain the subscript type based on the assignments, we still perform default type inference + // (to store the expression type and to report errors). + let slice_ty = self.infer_expression(slice, TypeContext::default()); + self.infer_subscript_expression_types(subscript, value_ty, slice_ty, *ctx); + return ty; + } + } + } + + let tuple_generic_alias = |db: &'db dyn Db, tuple: Option>| { + let tuple = tuple.unwrap_or_else(|| TupleType::homogeneous(db, Type::unknown())); + Type::from(tuple.to_class_type(db)) + }; + + match value_ty { + Type::ClassLiteral(class) => { + // HACK ALERT: If we are subscripting a generic class, short-circuit the rest of the + // subscript inference logic and treat this as an explicit specialization. + // TODO: Move this logic into a custom callable, and update `find_name_in_mro` to return + // this callable as the `__class_getitem__` method on `type`. That probably requires + // updating all of the subscript logic below to use custom callables for all of the _other_ + // special cases, too. + if class.is_tuple(self.db()) { + return tuple_generic_alias( + self.db(), + self.infer_tuple_type_expression(subscript), + ); + } else if class.is_known(self.db(), KnownClass::Type) { + let argument_ty = self.infer_type_expression(slice); + return Type::KnownInstance(KnownInstanceType::TypeGenericAlias( + InternedType::new(self.db(), argument_ty), + )); + } + + if let Some(generic_context) = class.generic_context(self.db()) + && let Some(class) = class.as_static() + { + return self.infer_explicit_class_specialization( + subscript, + value_ty, + class, + generic_context, + ); + } + } + Type::KnownInstance(KnownInstanceType::TypeAliasType(TypeAliasType::ManualPEP695( + _, + ))) => { + let slice_ty = self.infer_expression(slice, TypeContext::default()); + let mut variables = FxOrderSet::default(); + slice_ty.bind_and_find_all_legacy_typevars( + self.db(), + self.typevar_binding_context, + &mut variables, + ); + let generic_context = GenericContext::from_typevar_instances(self.db(), variables); + return Type::Dynamic(DynamicType::UnknownGeneric(generic_context)); + } + Type::KnownInstance(KnownInstanceType::TypeAliasType(type_alias)) => { + if let Some(generic_context) = type_alias.generic_context(self.db()) { + return self.infer_explicit_type_alias_type_specialization( + subscript, + value_ty, + type_alias, + generic_context, + ); + } + } + Type::SpecialForm(special_form) => match special_form { + SpecialFormType::Tuple => { + return tuple_generic_alias( + self.db(), + self.infer_tuple_type_expression(subscript), + ); + } + SpecialFormType::Literal => match self.infer_literal_parameter_type(slice) { + Ok(result) => { + return Type::KnownInstance(KnownInstanceType::Literal(InternedType::new( + self.db(), + result, + ))); + } + Err(nodes) => { + for node in nodes { + let Some(builder) = self.context.report_lint(&INVALID_TYPE_FORM, node) + else { + continue; + }; + builder.into_diagnostic( + "Type arguments for `Literal` must be `None`, \ + a literal value (int, bool, str, or bytes), or an enum member", + ); + } + return Type::unknown(); + } + }, + SpecialFormType::Annotated => { + let ast::Expr::Tuple(ast::ExprTuple { + elts: ref arguments, + .. + }) = **slice + else { + report_invalid_arguments_to_annotated(&self.context, subscript); + + return self.infer_expression(slice, TypeContext::default()); + }; + + if arguments.len() < 2 { + report_invalid_arguments_to_annotated(&self.context, subscript); + } + + let [type_expr, metadata @ ..] = &arguments[..] else { + for argument in arguments { + self.infer_expression(argument, TypeContext::default()); + } + self.store_expression_type(slice, Type::unknown()); + return Type::unknown(); + }; + + for element in metadata { + self.infer_expression(element, TypeContext::default()); + } + + let ty = self.infer_type_expression(type_expr); + + return Type::KnownInstance(KnownInstanceType::Annotated(InternedType::new( + self.db(), + ty, + ))); + } + SpecialFormType::Optional => { + let db = self.db(); + + if matches!(**slice, ast::Expr::Tuple(_)) + && let Some(builder) = + self.context.report_lint(&INVALID_TYPE_FORM, subscript) + { + builder.into_diagnostic(format_args!( + "`typing.Optional` requires exactly one argument" + )); + } + + let ty = self.infer_type_expression(slice); + + // `Optional[None]` is equivalent to `None`: + if ty.is_none(db) { + return ty; + } + + return Type::KnownInstance(KnownInstanceType::UnionType( + UnionTypeInstance::new( + db, + None, + Ok(UnionType::from_two_elements(db, ty, Type::none(db))), + ), + )); + } + SpecialFormType::Union => { + let db = self.db(); + + match **slice { + ast::Expr::Tuple(ref tuple) => { + let mut elements = tuple + .elts + .iter() + .map(|elt| self.infer_type_expression(elt)) + .peekable(); + + let is_empty = elements.peek().is_none(); + let union_type = Type::KnownInstance(KnownInstanceType::UnionType( + UnionTypeInstance::new( + db, + None, + Ok(UnionType::from_elements(db, elements)), + ), + )); + + if is_empty + && let Some(builder) = + self.context.report_lint(&INVALID_TYPE_FORM, subscript) + { + builder.into_diagnostic( + "`typing.Union` requires at least one type argument", + ); + } + + return union_type; + } + _ => { + return self.infer_expression(slice, TypeContext::default()); + } + } + } + SpecialFormType::Type => { + // Similar to the branch above that handles `type[…]`, handle `typing.Type[…]` + let argument_ty = self.infer_type_expression(slice); + return Type::KnownInstance(KnownInstanceType::TypeGenericAlias( + InternedType::new(self.db(), argument_ty), + )); + } + SpecialFormType::Callable => { + let arguments = if let ast::Expr::Tuple(tuple) = &*subscript.slice { + &*tuple.elts + } else { + std::slice::from_ref(&*subscript.slice) + }; + + // TODO: Remove this once we support Concatenate properly. This is necessary + // to avoid a lot of false positives downstream, because we can't represent the typevar- + // specialized `Callable` types yet. + let num_arguments = arguments.len(); + if num_arguments == 2 { + let first_arg = &arguments[0]; + let second_arg = &arguments[1]; + + if first_arg.is_subscript_expr() { + let first_arg_ty = + self.infer_expression(first_arg, TypeContext::default()); + if let Type::Dynamic(DynamicType::UnknownGeneric(generic_context)) = + first_arg_ty + { + let mut variables = generic_context + .variables(self.db()) + .collect::>(); + + let return_ty = + self.infer_expression(second_arg, TypeContext::default()); + return_ty.bind_and_find_all_legacy_typevars( + self.db(), + self.typevar_binding_context, + &mut variables, + ); + + let generic_context = + GenericContext::from_typevar_instances(self.db(), variables); + return Type::Dynamic(DynamicType::UnknownGeneric(generic_context)); + } + + if let Some(builder) = + self.context.report_lint(&INVALID_TYPE_FORM, subscript) + { + builder.into_diagnostic(format_args!( + "The first argument to `Callable` must be either a list of types, \ + ParamSpec, Concatenate, or `...`", + )); + } + return Type::KnownInstance(KnownInstanceType::Callable( + CallableType::unknown(self.db()), + )); + } + } + + let callable = self + .infer_callable_type(subscript) + .as_callable() + .expect("always returns Type::Callable"); + + return Type::KnownInstance(KnownInstanceType::Callable(callable)); + } + SpecialFormType::LegacyStdlibAlias(alias) => { + let AliasSpec { + class, + expected_argument_number, + } = alias.alias_spec(); + + let args = if let ast::Expr::Tuple(t) = &**slice { + &*t.elts + } else { + std::slice::from_ref(&**slice) + }; + + if args.len() != expected_argument_number { + if let Some(builder) = + self.context.report_lint(&INVALID_TYPE_FORM, subscript) + { + let noun = if expected_argument_number == 1 { + "argument" + } else { + "arguments" + }; + builder.into_diagnostic(format_args!( + "`typing.{name}` requires exactly \ + {expected_argument_number} {noun}, got {got}", + name = special_form.name(), + got = args.len() + )); + } + } + + let arg_types: Vec<_> = args + .iter() + .map(|arg| self.infer_type_expression(arg)) + .collect(); + + return class + .to_specialized_class_type(self.db(), arg_types) + .map(Type::from) + .unwrap_or_else(Type::unknown); + } + _ => {} + }, + + Type::KnownInstance( + KnownInstanceType::UnionType(_) + | KnownInstanceType::Annotated(_) + | KnownInstanceType::Callable(_) + | KnownInstanceType::TypeGenericAlias(_), + ) => { + return self.infer_explicit_type_alias_specialization(subscript, value_ty, false); + } + Type::Dynamic(DynamicType::Unknown) => { + let slice_ty = self.infer_expression(slice, TypeContext::default()); + let mut variables = FxOrderSet::default(); + slice_ty.bind_and_find_all_legacy_typevars( + self.db(), + self.typevar_binding_context, + &mut variables, + ); + let generic_context = GenericContext::from_typevar_instances(self.db(), variables); + return Type::Dynamic(DynamicType::UnknownGeneric(generic_context)); + } + _ => {} + } + + let slice_ty = self.infer_expression(slice, TypeContext::default()); + let result_ty = self.infer_subscript_expression_types(subscript, value_ty, slice_ty, *ctx); + self.narrow_expr_with_applicable_constraints(subscript, result_ty, &constraint_keys) + } + + pub(super) fn infer_explicit_class_specialization( + &mut self, + subscript: &ast::ExprSubscript, + value_ty: Type<'db>, + generic_class: StaticClassLiteral<'db>, + generic_context: GenericContext<'db>, + ) -> Type<'db> { + let db = self.db(); + let specialize = &|types: &[Option>]| { + Type::from(generic_class.apply_specialization(db, |_| { + generic_context.specialize_partial(db, types.iter().copied()) + })) + }; + + self.infer_explicit_callable_specialization( + subscript, + value_ty, + generic_context, + specialize, + ) + } + + pub(super) fn infer_explicit_type_alias_type_specialization( + &mut self, + subscript: &ast::ExprSubscript, + value_ty: Type<'db>, + generic_type_alias: TypeAliasType<'db>, + generic_context: GenericContext<'db>, + ) -> Type<'db> { + let db = self.db(); + let specialize = &|types: &[Option>]| { + let type_alias = generic_type_alias.apply_specialization(db, |_| { + generic_context.specialize_partial(db, types.iter().copied()) + }); + + Type::KnownInstance(KnownInstanceType::TypeAliasType(type_alias)) + }; + + self.infer_explicit_callable_specialization( + subscript, + value_ty, + generic_context, + specialize, + ) + } + + pub(super) fn infer_explicit_callable_specialization( + &mut self, + subscript: &ast::ExprSubscript, + value_ty: Type<'db>, + generic_context: GenericContext<'db>, + specialize: &dyn Fn(&[Option>]) -> Type<'db>, + ) -> Type<'db> { + enum ExplicitSpecializationError { + InvalidParamSpec, + UnsatisfiedBound, + UnsatisfiedConstraints, + /// These two errors override the errors above, causing all specializations to be `Unknown`. + MissingTypeVars, + TooManyArguments, + /// This error overrides the errors above, causing the type itself to be `Unknown`. + NonGeneric, + } + + fn add_typevar_definition<'db>( + db: &'db dyn Db, + diagnostic: &mut Diagnostic, + typevar: BoundTypeVarInstance<'db>, + ) { + let Some(definition) = typevar.typevar(db).definition(db) else { + return; + }; + let file = definition.file(db); + let module = parsed_module(db, file).load(db); + let range = definition.focus_range(db, &module).range(); + diagnostic.annotate( + Annotation::secondary(Span::from(file).with_range(range)) + .message("Type variable defined here"), + ); + } + + let db = self.db(); + let constraints = ConstraintSetBuilder::new(); + let slice_node = subscript.slice.as_ref(); + + let exactly_one_paramspec = generic_context.exactly_one_paramspec(db); + let (type_arguments, store_inferred_type_arguments) = match slice_node { + ast::Expr::Tuple(tuple) => { + if exactly_one_paramspec && !tuple.elts.is_empty() { + (std::slice::from_ref(slice_node), false) + } else { + (tuple.elts.as_slice(), true) + } + } + _ => (std::slice::from_ref(slice_node), false), + }; + let mut inferred_type_arguments = Vec::with_capacity(type_arguments.len()); + + let typevars = generic_context.variables(db); + let typevars_len = typevars.len(); + + let mut specialization_types = Vec::with_capacity(typevars_len); + let mut typevar_with_defaults = 0; + let mut missing_typevars = vec![]; + let mut first_excess_type_argument_index = None; + + // Helper to get the AST node corresponding to the type argument at `index`. + let get_node = |index: usize| -> ast::AnyNodeRef<'_> { + match slice_node { + ast::Expr::Tuple(ast::ExprTuple { elts, .. }) if !exactly_one_paramspec => elts + .get(index) + .expect("type argument index should not be out of range") + .into(), + _ => slice_node.into(), + } + }; + + let mut error: Option = None; + + for (index, item) in typevars.zip_longest(type_arguments.iter()).enumerate() { + match item { + EitherOrBoth::Both(typevar, expr) => { + if typevar.default_type(db).is_some() { + typevar_with_defaults += 1; + } + + let provided_type = if typevar.is_paramspec(db) { + match self.infer_paramspec_explicit_specialization_value( + expr, + exactly_one_paramspec, + ) { + Ok(paramspec_value) => paramspec_value, + Err(()) => { + error = Some(ExplicitSpecializationError::InvalidParamSpec); + Type::paramspec_value_callable(db, Parameters::unknown()) + } + } + } else { + self.infer_type_expression(expr) + }; + + inferred_type_arguments.push(provided_type); + + // TODO consider just accepting the given specialization without checking + // against bounds/constraints, but recording the expression for deferred + // checking at end of scope. This would avoid a lot of cycles caused by eagerly + // doing assignment checks here. + match typevar.typevar(db).bound_or_constraints(db) { + Some(TypeVarBoundOrConstraints::UpperBound(bound)) => { + if provided_type + .when_assignable_to( + db, + bound, + &constraints, + InferableTypeVars::None, + ) + .is_never_satisfied(db) + { + let node = get_node(index); + if let Some(builder) = + self.context.report_lint(&INVALID_TYPE_ARGUMENTS, node) + { + let mut diagnostic = builder.into_diagnostic(format_args!( + "Type `{}` is not assignable to upper bound `{}` \ + of type variable `{}`", + provided_type.display(db), + bound.display(db), + typevar.identity(db).display(db), + )); + add_typevar_definition(db, &mut diagnostic, typevar); + } + error = Some(ExplicitSpecializationError::UnsatisfiedBound); + specialization_types.push(Some(Type::unknown())); + } else { + specialization_types.push(Some(provided_type)); + } + } + Some(TypeVarBoundOrConstraints::Constraints(typevar_constraints)) => { + // TODO: this is wrong, the given specialization needs to be assignable + // to _at least one_ of the individual constraints, not to the union of + // all of them. `int | str` is not a valid specialization of a typevar + // constrained to `(int, str)`. + if provided_type + .when_assignable_to( + db, + typevar_constraints.as_type(db), + &constraints, + InferableTypeVars::None, + ) + .is_never_satisfied(db) + { + let node = get_node(index); + if let Some(builder) = + self.context.report_lint(&INVALID_TYPE_ARGUMENTS, node) + { + let mut diagnostic = builder.into_diagnostic(format_args!( + "Type `{}` does not satisfy constraints `{}` \ + of type variable `{}`", + provided_type.display(db), + typevar_constraints + .elements(db) + .iter() + .map(|c| c.display(db)) + .format("`, `"), + typevar.identity(db).display(db), + )); + add_typevar_definition(db, &mut diagnostic, typevar); + } + error = Some(ExplicitSpecializationError::UnsatisfiedConstraints); + specialization_types.push(Some(Type::unknown())); + } else { + specialization_types.push(Some(provided_type)); + } + } + None => { + specialization_types.push(Some(provided_type)); + } + } + } + EitherOrBoth::Left(typevar) => { + if typevar.default_type(db).is_none() { + // This is an error case, so no need to push into the specialization types. + missing_typevars.push(typevar); + } else { + typevar_with_defaults += 1; + specialization_types.push(None); + } + } + EitherOrBoth::Right(expr) => { + inferred_type_arguments.push(self.infer_type_expression(expr)); + first_excess_type_argument_index.get_or_insert(index); + } + } + } + + if !missing_typevars.is_empty() { + if let Some(builder) = self.context.report_lint(&INVALID_TYPE_ARGUMENTS, subscript) { + let description = CallableDescription::new(db, value_ty); + let s = if missing_typevars.len() > 1 { "s" } else { "" }; + builder.into_diagnostic(format_args!( + "No type argument{s} provided for required type variable{s} `{}`{}", + missing_typevars + .iter() + .map(|tv| tv.typevar(db).name(db)) + .format("`, `"), + if let Some(CallableDescription { kind, name }) = description { + format!(" of {kind} `{name}`") + } else { + String::new() + } + )); + } + error = Some(ExplicitSpecializationError::MissingTypeVars); + } + + if let Some(first_excess_type_argument_index) = first_excess_type_argument_index { + if let Type::GenericAlias(alias) = value_ty + && let spec = alias.specialization(self.db()) + && spec + .types(self.db()) + .contains(&Type::Dynamic(DynamicType::TodoTypeVarTuple)) + { + // Avoid false-positive errors when specializing a class + // that's generic over a legacy TypeVarTuple + } else if typevars_len == 0 { + // Type parameter list cannot be empty, so if we reach here, `value_ty` is not a generic type. + if let Some(builder) = self.context.report_lint(&NOT_SUBSCRIPTABLE, subscript) { + let mut diagnostic = builder.into_diagnostic(format_args!( + "Cannot subscript non-generic type `{}`", + value_ty.display(db) + )); + if match value_ty { + Type::GenericAlias(_) => true, + Type::KnownInstance(KnownInstanceType::UnionType(union)) => union + .value_expression_types(db) + .is_ok_and(|mut tys| tys.any(|ty| ty.is_generic_alias())), + _ => false, + } { + diagnostic.annotate( + self.context + .secondary(&*subscript.value) + .message("Type is already specialized"), + ); + } + } + error = Some(ExplicitSpecializationError::NonGeneric); + } else { + let node = get_node(first_excess_type_argument_index); + if let Some(builder) = self.context.report_lint(&INVALID_TYPE_ARGUMENTS, node) { + let description = CallableDescription::new(db, value_ty); + builder.into_diagnostic(format_args!( + "Too many type arguments{}: expected {}, got {}", + if let Some(CallableDescription { kind, name }) = description { + format!(" to {kind} `{name}`") + } else { + String::new() + }, + if typevar_with_defaults == 0 { + format!("{typevars_len}") + } else { + format!( + "between {} and {}", + typevars_len - typevar_with_defaults, + typevars_len + ) + }, + type_arguments.len(), + )); + } + error = Some(ExplicitSpecializationError::TooManyArguments); + } + } + + if store_inferred_type_arguments { + self.store_expression_type( + slice_node, + Type::heterogeneous_tuple(db, inferred_type_arguments), + ); + } + + match error { + Some(ExplicitSpecializationError::NonGeneric) => Type::unknown(), + Some( + ExplicitSpecializationError::MissingTypeVars + | ExplicitSpecializationError::TooManyArguments, + ) => { + let unknowns = generic_context + .variables(self.db()) + .map(|typevar| { + Some(if typevar.is_paramspec(db) { + Type::paramspec_value_callable(db, Parameters::unknown()) + } else { + Type::unknown() + }) + }) + .collect::>(); + specialize(&unknowns) + } + Some( + ExplicitSpecializationError::UnsatisfiedBound + | ExplicitSpecializationError::UnsatisfiedConstraints + | ExplicitSpecializationError::InvalidParamSpec, + ) + | None => specialize(&specialization_types), + } + } + + pub(super) fn infer_subscript_expression_types( + &self, + subscript: &ast::ExprSubscript, + value_ty: Type<'db>, + slice_ty: Type<'db>, + expr_context: ExprContext, + ) -> Type<'db> { + let db = self.db(); + + // Special typing forms for which subscriptions are context-dependent are parsed here, + // outside of `Type::subscript`, which is a pure function that doesn't depend on the + // semantic index or any context-dependent state. + let subscript_result = match value_ty { + Type::SpecialForm(SpecialFormType::Generic) => infer_legacy_generic_subscript( + db, + self.index, + self.scope().file_scope_id(db), + self.typevar_binding_context, + slice_ty, + LegacyGenericOrigin::Generic, + KnownInstanceType::SubscriptedGeneric, + ), + Type::SpecialForm(SpecialFormType::Protocol) => infer_legacy_generic_subscript( + db, + self.index, + self.scope().file_scope_id(db), + self.typevar_binding_context, + slice_ty, + LegacyGenericOrigin::Protocol, + KnownInstanceType::SubscriptedProtocol, + ), + Type::SpecialForm(SpecialFormType::Concatenate) => { + // TODO: Add proper support for `Concatenate` + let mut variables = FxOrderSet::default(); + slice_ty.bind_and_find_all_legacy_typevars( + db, + self.typevar_binding_context, + &mut variables, + ); + let generic_context = GenericContext::from_typevar_instances(db, variables); + Ok(Type::Dynamic(DynamicType::UnknownGeneric(generic_context))) + } + _ => value_ty.subscript(self.db(), slice_ty, expr_context), + }; + + subscript_result.unwrap_or_else(|e| { + e.report_diagnostics(&self.context, subscript); + e.result_type() + }) + } + + pub(super) fn infer_slice_expression(&mut self, slice: &ast::ExprSlice) -> Type<'db> { + enum SliceArg<'db> { + Arg(Type<'db>), + Unsupported, + } + + let ast::ExprSlice { + range: _, + node_index: _, + lower, + upper, + step, + } = slice; + + let ty_lower = self.infer_optional_expression(lower.as_deref(), TypeContext::default()); + let ty_upper = self.infer_optional_expression(upper.as_deref(), TypeContext::default()); + let ty_step = self.infer_optional_expression(step.as_deref(), TypeContext::default()); + + let type_to_slice_argument = |ty: Option>| match ty { + Some(ty @ Type::LiteralValue(literal)) if literal.is_int() || literal.is_bool() => { + SliceArg::Arg(ty) + } + Some(ty @ Type::NominalInstance(instance)) + if instance.has_known_class(self.db(), KnownClass::NoneType) => + { + SliceArg::Arg(ty) + } + None => SliceArg::Arg(Type::none(self.db())), + _ => SliceArg::Unsupported, + }; + + match ( + type_to_slice_argument(ty_lower), + type_to_slice_argument(ty_upper), + type_to_slice_argument(ty_step), + ) { + (SliceArg::Arg(lower), SliceArg::Arg(upper), SliceArg::Arg(step)) => { + KnownClass::Slice.to_specialized_instance(self.db(), &[lower, upper, step]) + } + _ => KnownClass::Slice.to_instance(self.db()), + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum LegacyGenericContextError<'db> { + /// It's invalid to subscript `Generic` or `Protocol` with this type. + InvalidArgument(Type<'db>), + /// It's invalid to subscript `Generic` or `Protocol` with a variadic tuple type. + /// We should emit a diagnostic for this, but we don't yet. + VariadicTupleArguments, + /// It's valid to subscribe `Generic` or `Protocol` with this type, + /// but the type is not yet supported. + NotYetSupported, + /// A duplicate typevar was provided. + DuplicateTypevar(&'db str), + /// A `TypeVarTuple` was provided but not unpacked. + TypeVarTupleMustBeUnpacked, +} + +impl<'db> LegacyGenericContextError<'db> { + const fn into_type(self) -> Type<'db> { + match self { + LegacyGenericContextError::InvalidArgument(_) + | LegacyGenericContextError::VariadicTupleArguments + | LegacyGenericContextError::DuplicateTypevar(_) + | LegacyGenericContextError::TypeVarTupleMustBeUnpacked => Type::unknown(), + LegacyGenericContextError::NotYetSupported => { + todo_type!("ParamSpecs and TypeVarTuples") + } + } + } +} + +/// Validate the type arguments to `Generic[...]` or `Protocol[...]`, returning +/// either the resulting [`GenericContext`] or a [`SubscriptError`]. +fn infer_legacy_generic_subscript<'db>( + db: &'db dyn Db, + index: &'db SemanticIndex<'db>, + file_scope_id: FileScopeId, + typevar_binding_context: Option>, + slice_ty: Type<'db>, + origin: LegacyGenericOrigin, + wrap_ok: impl FnOnce(GenericContext<'db>) -> KnownInstanceType<'db>, +) -> Result, SubscriptError<'db>> { + match legacy_generic_class_context(db, index, file_scope_id, typevar_binding_context, slice_ty) + { + Ok(context) => Ok(Type::KnownInstance(wrap_ok(context))), + Err(LegacyGenericContextError::InvalidArgument(argument_ty)) => Err(SubscriptError::new( + Type::unknown(), + SubscriptErrorKind::InvalidLegacyGenericArgument { + origin, + argument_ty, + }, + )), + Err(LegacyGenericContextError::DuplicateTypevar(typevar_name)) => Err(SubscriptError::new( + Type::unknown(), + SubscriptErrorKind::DuplicateTypevar { + origin, + typevar_name, + }, + )), + Err(LegacyGenericContextError::TypeVarTupleMustBeUnpacked) => Err(SubscriptError::new( + Type::unknown(), + SubscriptErrorKind::TypeVarTupleNotUnpacked { origin }, + )), + Err( + error @ (LegacyGenericContextError::NotYetSupported + | LegacyGenericContextError::VariadicTupleArguments), + ) => Ok(error.into_type()), + } +} + +/// Parse the type arguments to `Generic[...]` or `Protocol[...]` and validate +/// that each argument is a type variable. +fn legacy_generic_class_context<'db>( + db: &'db dyn Db, + index: &'db SemanticIndex<'db>, + file_scope_id: FileScopeId, + typevar_binding_context: Option>, + typevars: Type<'db>, +) -> Result, LegacyGenericContextError<'db>> { + let typevars_class_tuple_spec = typevars.exact_tuple_instance_spec(db); + + let typevars = if let Some(tuple_spec) = typevars_class_tuple_spec.as_deref() { + match tuple_spec { + Tuple::Fixed(typevars) => typevars.elements_slice(), + Tuple::Variable(_) => { + return Err(LegacyGenericContextError::VariadicTupleArguments); + } + } + } else { + std::slice::from_ref(&typevars) + }; + + let mut validated_typevars = FxOrderSet::default(); + for ty in typevars { + let argument_ty = *ty; + if let Type::KnownInstance(KnownInstanceType::TypeVar(typevar)) = argument_ty { + let bound = bind_typevar(db, index, file_scope_id, typevar_binding_context, typevar) + .ok_or(LegacyGenericContextError::InvalidArgument(argument_ty))?; + if !validated_typevars.insert(bound) { + return Err(LegacyGenericContextError::DuplicateTypevar( + typevar.name(db), + )); + } + } else if let Type::NominalInstance(instance) = argument_ty + && instance.has_known_class(db, KnownClass::TypeVarTuple) + { + return Err(LegacyGenericContextError::TypeVarTupleMustBeUnpacked); + } else if any_over_type(db, argument_ty, true, |inner_ty| match inner_ty { + Type::Dynamic(DynamicType::TodoUnpack | DynamicType::TodoStarredExpression) => true, + Type::NominalInstance(nominal) => nominal.has_known_class(db, KnownClass::TypeVarTuple), + _ => false, + }) { + return Err(LegacyGenericContextError::NotYetSupported); + } else { + return Err(LegacyGenericContextError::InvalidArgument(argument_ty)); + } + } + Ok(GenericContext::from_typevar_instances( + db, + validated_typevars, + )) +} From 6f4872ea0e91c7269a6c108ecadf44697ff4a833 Mon Sep 17 00:00:00 2001 From: Alex Waygood Date: Mon, 2 Mar 2026 12:22:48 +0000 Subject: [PATCH 2/2] cleanups --- .../src/types/infer/builder/subscript.rs | 232 ++++++++---------- 1 file changed, 107 insertions(+), 125 deletions(-) diff --git a/crates/ty_python_semantic/src/types/infer/builder/subscript.rs b/crates/ty_python_semantic/src/types/infer/builder/subscript.rs index 474b668d2e0bd6..acdc9da74097f1 100644 --- a/crates/ty_python_semantic/src/types/infer/builder/subscript.rs +++ b/crates/ty_python_semantic/src/types/infer/builder/subscript.rs @@ -81,6 +81,8 @@ impl<'db> TypeInferenceBuilder<'db, '_> { value_ty: Type<'db>, subscript: &ast::ExprSubscript, ) -> Type<'db> { + let db = self.db(); + let ast::ExprSubscript { range: _, node_index: _, @@ -127,19 +129,16 @@ impl<'db> TypeInferenceBuilder<'db, '_> { // this callable as the `__class_getitem__` method on `type`. That probably requires // updating all of the subscript logic below to use custom callables for all of the _other_ // special cases, too. - if class.is_tuple(self.db()) { - return tuple_generic_alias( - self.db(), - self.infer_tuple_type_expression(subscript), - ); - } else if class.is_known(self.db(), KnownClass::Type) { + if class.is_tuple(db) { + return tuple_generic_alias(db, self.infer_tuple_type_expression(subscript)); + } else if class.is_known(db, KnownClass::Type) { let argument_ty = self.infer_type_expression(slice); return Type::KnownInstance(KnownInstanceType::TypeGenericAlias( - InternedType::new(self.db(), argument_ty), + InternedType::new(db, argument_ty), )); } - if let Some(generic_context) = class.generic_context(self.db()) + if let Some(generic_context) = class.generic_context(db) && let Some(class) = class.as_static() { return self.infer_explicit_class_specialization( @@ -156,15 +155,15 @@ impl<'db> TypeInferenceBuilder<'db, '_> { let slice_ty = self.infer_expression(slice, TypeContext::default()); let mut variables = FxOrderSet::default(); slice_ty.bind_and_find_all_legacy_typevars( - self.db(), + db, self.typevar_binding_context, &mut variables, ); - let generic_context = GenericContext::from_typevar_instances(self.db(), variables); + let generic_context = GenericContext::from_typevar_instances(db, variables); return Type::Dynamic(DynamicType::UnknownGeneric(generic_context)); } Type::KnownInstance(KnownInstanceType::TypeAliasType(type_alias)) => { - if let Some(generic_context) = type_alias.generic_context(self.db()) { + if let Some(generic_context) = type_alias.generic_context(db) { return self.infer_explicit_type_alias_type_specialization( subscript, value_ty, @@ -175,16 +174,12 @@ impl<'db> TypeInferenceBuilder<'db, '_> { } Type::SpecialForm(special_form) => match special_form { SpecialFormType::Tuple => { - return tuple_generic_alias( - self.db(), - self.infer_tuple_type_expression(subscript), - ); + return tuple_generic_alias(db, self.infer_tuple_type_expression(subscript)); } SpecialFormType::Literal => match self.infer_literal_parameter_type(slice) { Ok(result) => { return Type::KnownInstance(KnownInstanceType::Literal(InternedType::new( - self.db(), - result, + db, result, ))); } Err(nodes) => { @@ -231,13 +226,10 @@ impl<'db> TypeInferenceBuilder<'db, '_> { let ty = self.infer_type_expression(type_expr); return Type::KnownInstance(KnownInstanceType::Annotated(InternedType::new( - self.db(), - ty, + db, ty, ))); } SpecialFormType::Optional => { - let db = self.db(); - if matches!(**slice, ast::Expr::Tuple(_)) && let Some(builder) = self.context.report_lint(&INVALID_TYPE_FORM, subscript) @@ -262,47 +254,43 @@ impl<'db> TypeInferenceBuilder<'db, '_> { ), )); } - SpecialFormType::Union => { - let db = self.db(); - - match **slice { - ast::Expr::Tuple(ref tuple) => { - let mut elements = tuple - .elts - .iter() - .map(|elt| self.infer_type_expression(elt)) - .peekable(); - - let is_empty = elements.peek().is_none(); - let union_type = Type::KnownInstance(KnownInstanceType::UnionType( - UnionTypeInstance::new( - db, - None, - Ok(UnionType::from_elements(db, elements)), - ), - )); - - if is_empty - && let Some(builder) = - self.context.report_lint(&INVALID_TYPE_FORM, subscript) - { - builder.into_diagnostic( - "`typing.Union` requires at least one type argument", - ); - } + SpecialFormType::Union => match **slice { + ast::Expr::Tuple(ref tuple) => { + let mut elements = tuple + .elts + .iter() + .map(|elt| self.infer_type_expression(elt)) + .peekable(); + + let is_empty = elements.peek().is_none(); + let union_type = Type::KnownInstance(KnownInstanceType::UnionType( + UnionTypeInstance::new( + db, + None, + Ok(UnionType::from_elements(db, elements)), + ), + )); - return union_type; - } - _ => { - return self.infer_expression(slice, TypeContext::default()); + if is_empty + && let Some(builder) = + self.context.report_lint(&INVALID_TYPE_FORM, subscript) + { + builder.into_diagnostic( + "`typing.Union` requires at least one type argument", + ); } + + return union_type; } - } + _ => { + return self.infer_expression(slice, TypeContext::default()); + } + }, SpecialFormType::Type => { // Similar to the branch above that handles `type[…]`, handle `typing.Type[…]` let argument_ty = self.infer_type_expression(slice); return Type::KnownInstance(KnownInstanceType::TypeGenericAlias( - InternedType::new(self.db(), argument_ty), + InternedType::new(db, argument_ty), )); } SpecialFormType::Callable => { @@ -315,46 +303,40 @@ impl<'db> TypeInferenceBuilder<'db, '_> { // TODO: Remove this once we support Concatenate properly. This is necessary // to avoid a lot of false positives downstream, because we can't represent the typevar- // specialized `Callable` types yet. - let num_arguments = arguments.len(); - if num_arguments == 2 { - let first_arg = &arguments[0]; - let second_arg = &arguments[1]; - - if first_arg.is_subscript_expr() { - let first_arg_ty = - self.infer_expression(first_arg, TypeContext::default()); - if let Type::Dynamic(DynamicType::UnknownGeneric(generic_context)) = - first_arg_ty - { - let mut variables = generic_context - .variables(self.db()) - .collect::>(); - - let return_ty = - self.infer_expression(second_arg, TypeContext::default()); - return_ty.bind_and_find_all_legacy_typevars( - self.db(), - self.typevar_binding_context, - &mut variables, - ); - - let generic_context = - GenericContext::from_typevar_instances(self.db(), variables); - return Type::Dynamic(DynamicType::UnknownGeneric(generic_context)); - } + if let [first_arg, second_arg] = arguments + && first_arg.is_subscript_expr() + { + let first_arg_ty = self.infer_expression(first_arg, TypeContext::default()); + if let Type::Dynamic(DynamicType::UnknownGeneric(generic_context)) = + first_arg_ty + { + let mut variables = + generic_context.variables(db).collect::>(); + + let return_ty = + self.infer_expression(second_arg, TypeContext::default()); + return_ty.bind_and_find_all_legacy_typevars( + db, + self.typevar_binding_context, + &mut variables, + ); - if let Some(builder) = - self.context.report_lint(&INVALID_TYPE_FORM, subscript) - { - builder.into_diagnostic(format_args!( + let generic_context = + GenericContext::from_typevar_instances(db, variables); + return Type::Dynamic(DynamicType::UnknownGeneric(generic_context)); + } + + if let Some(builder) = + self.context.report_lint(&INVALID_TYPE_FORM, subscript) + { + builder.into_diagnostic(format_args!( "The first argument to `Callable` must be either a list of types, \ ParamSpec, Concatenate, or `...`", )); - } - return Type::KnownInstance(KnownInstanceType::Callable( - CallableType::unknown(self.db()), - )); } + return Type::KnownInstance(KnownInstanceType::Callable( + CallableType::unknown(db), + )); } let callable = self @@ -376,22 +358,21 @@ impl<'db> TypeInferenceBuilder<'db, '_> { std::slice::from_ref(&**slice) }; - if args.len() != expected_argument_number { - if let Some(builder) = + if args.len() != expected_argument_number + && let Some(builder) = self.context.report_lint(&INVALID_TYPE_FORM, subscript) - { - let noun = if expected_argument_number == 1 { - "argument" - } else { - "arguments" - }; - builder.into_diagnostic(format_args!( - "`typing.{name}` requires exactly \ + { + let noun = if expected_argument_number == 1 { + "argument" + } else { + "arguments" + }; + builder.into_diagnostic(format_args!( + "`typing.{name}` requires exactly \ {expected_argument_number} {noun}, got {got}", - name = special_form.name(), - got = args.len() - )); - } + name = special_form.name(), + got = args.len() + )); } let arg_types: Vec<_> = args @@ -400,7 +381,7 @@ impl<'db> TypeInferenceBuilder<'db, '_> { .collect(); return class - .to_specialized_class_type(self.db(), arg_types) + .to_specialized_class_type(db, arg_types) .map(Type::from) .unwrap_or_else(Type::unknown); } @@ -419,11 +400,11 @@ impl<'db> TypeInferenceBuilder<'db, '_> { let slice_ty = self.infer_expression(slice, TypeContext::default()); let mut variables = FxOrderSet::default(); slice_ty.bind_and_find_all_legacy_typevars( - self.db(), + db, self.typevar_binding_context, &mut variables, ); - let generic_context = GenericContext::from_typevar_instances(self.db(), variables); + let generic_context = GenericContext::from_typevar_instances(db, variables); return Type::Dynamic(DynamicType::UnknownGeneric(generic_context)); } _ => {} @@ -561,16 +542,14 @@ impl<'db> TypeInferenceBuilder<'db, '_> { } let provided_type = if typevar.is_paramspec(db) { - match self.infer_paramspec_explicit_specialization_value( + self.infer_paramspec_explicit_specialization_value( expr, exactly_one_paramspec, - ) { - Ok(paramspec_value) => paramspec_value, - Err(()) => { - error = Some(ExplicitSpecializationError::InvalidParamSpec); - Type::paramspec_value_callable(db, Parameters::unknown()) - } - } + ) + .unwrap_or_else(|()| { + error = Some(ExplicitSpecializationError::InvalidParamSpec); + Type::paramspec_value_callable(db, Parameters::unknown()) + }) } else { self.infer_type_expression(expr) }; @@ -691,9 +670,9 @@ impl<'db> TypeInferenceBuilder<'db, '_> { if let Some(first_excess_type_argument_index) = first_excess_type_argument_index { if let Type::GenericAlias(alias) = value_ty - && let spec = alias.specialization(self.db()) + && let spec = alias.specialization(db) && spec - .types(self.db()) + .types(db) .contains(&Type::Dynamic(DynamicType::TodoTypeVarTuple)) { // Avoid false-positive errors when specializing a class @@ -705,13 +684,14 @@ impl<'db> TypeInferenceBuilder<'db, '_> { "Cannot subscript non-generic type `{}`", value_ty.display(db) )); - if match value_ty { + let already_specialized = match value_ty { Type::GenericAlias(_) => true, Type::KnownInstance(KnownInstanceType::UnionType(union)) => union .value_expression_types(db) .is_ok_and(|mut tys| tys.any(|ty| ty.is_generic_alias())), _ => false, - } { + }; + if already_specialized { diagnostic.annotate( self.context .secondary(&*subscript.value) @@ -761,7 +741,7 @@ impl<'db> TypeInferenceBuilder<'db, '_> { | ExplicitSpecializationError::TooManyArguments, ) => { let unknowns = generic_context - .variables(self.db()) + .variables(db) .map(|typevar| { Some(if typevar.is_paramspec(db) { Type::paramspec_value_callable(db, Parameters::unknown()) @@ -823,7 +803,7 @@ impl<'db> TypeInferenceBuilder<'db, '_> { let generic_context = GenericContext::from_typevar_instances(db, variables); Ok(Type::Dynamic(DynamicType::UnknownGeneric(generic_context))) } - _ => value_ty.subscript(self.db(), slice_ty, expr_context), + _ => value_ty.subscript(db, slice_ty, expr_context), }; subscript_result.unwrap_or_else(|e| { @@ -838,6 +818,8 @@ impl<'db> TypeInferenceBuilder<'db, '_> { Unsupported, } + let db = self.db(); + let ast::ExprSlice { range: _, node_index: _, @@ -855,11 +837,11 @@ impl<'db> TypeInferenceBuilder<'db, '_> { SliceArg::Arg(ty) } Some(ty @ Type::NominalInstance(instance)) - if instance.has_known_class(self.db(), KnownClass::NoneType) => + if instance.has_known_class(db, KnownClass::NoneType) => { SliceArg::Arg(ty) } - None => SliceArg::Arg(Type::none(self.db())), + None => SliceArg::Arg(Type::none(db)), _ => SliceArg::Unsupported, }; @@ -869,9 +851,9 @@ impl<'db> TypeInferenceBuilder<'db, '_> { type_to_slice_argument(ty_step), ) { (SliceArg::Arg(lower), SliceArg::Arg(upper), SliceArg::Arg(step)) => { - KnownClass::Slice.to_specialized_instance(self.db(), &[lower, upper, step]) + KnownClass::Slice.to_specialized_instance(db, &[lower, upper, step]) } - _ => KnownClass::Slice.to_instance(self.db()), + _ => KnownClass::Slice.to_instance(db), } } }