From 96d77e3301b6e7e2024cd0bdcde504e642396122 Mon Sep 17 00:00:00 2001 From: Douglas Creager Date: Fri, 5 Dec 2025 14:18:21 -0500 Subject: [PATCH] pass in params when possible --- .../ty_python_semantic/src/semantic_index.rs | 5 + .../ty_python_semantic/src/types/function.rs | 17 +- crates/ty_python_semantic/src/types/infer.rs | 23 +-- .../src/types/infer/builder.rs | 22 ++- .../src/types/signatures.rs | 164 +++++------------- 5 files changed, 87 insertions(+), 144 deletions(-) diff --git a/crates/ty_python_semantic/src/semantic_index.rs b/crates/ty_python_semantic/src/semantic_index.rs index f4ab765f08..cbb5e9229f 100644 --- a/crates/ty_python_semantic/src/semantic_index.rs +++ b/crates/ty_python_semantic/src/semantic_index.rs @@ -525,6 +525,11 @@ impl<'db> SemanticIndex<'db> { self.scopes_by_node[&node.node_key()] } + #[track_caller] + pub(crate) fn try_node_scope(&self, node: NodeWithScopeRef) -> Option { + self.scopes_by_node.get(&node.node_key()).copied() + } + /// Checks if there is an import of `__future__.annotations` in the global scope, which affects /// the logic for type inference. pub(super) fn has_future_annotations(&self) -> bool { diff --git a/crates/ty_python_semantic/src/types/function.rs b/crates/ty_python_semantic/src/types/function.rs index 8d83b626c5..8089a74776 100644 --- a/crates/ty_python_semantic/src/types/function.rs +++ b/crates/ty_python_semantic/src/types/function.rs @@ -77,7 +77,7 @@ use crate::types::generics::{GenericContext, InferableTypeVars, typing_self}; use crate::types::infer::nearest_enclosing_class; use crate::types::list_members::all_members; use crate::types::narrow::ClassInfoConstraintFunction; -use crate::types::signatures::{CallableSignature, Signature}; +use crate::types::signatures::{CallableSignature, Parameter, Signature}; use crate::types::visitor::any_over_type; use crate::types::{ ApplyTypeMappingVisitor, BoundMethodType, BoundTypeVarInstance, CallableType, ClassBase, @@ -196,6 +196,12 @@ pub(crate) fn is_implicit_classmethod(function_name: &str) -> bool { matches!(function_name, "__init_subclass__" | "__class_getitem__") } +#[derive(Clone, Debug, Eq, Hash, PartialEq, get_size2::GetSize, salsa::Update)] +pub struct InferredFunctionSignature<'db> { + pub(crate) parameters: Vec>, + pub(crate) return_type: Option>, +} + /// Representation of a function definition in the AST: either a non-generic function, or a generic /// function that has not been specialized. /// @@ -228,6 +234,11 @@ pub struct OverloadLiteral<'db> { /// The arguments to `dataclass_transformer`, if this function was annotated /// with `@dataclass_transformer(...)`. pub(crate) dataclass_transformer_params: Option>, + + inferred_signature: Option>, + + #[returns(ref)] + inferred_defaults: Vec>>, } // The Salsa heap is tracked separately. @@ -248,6 +259,8 @@ impl<'db> OverloadLiteral<'db> { self.decorators(db), self.deprecated(db), Some(params), + self.inferred_signature(db), + self.inferred_defaults(db), ) } @@ -505,6 +518,8 @@ impl<'db> OverloadLiteral<'db> { pep695_ctx, definition, function_stmt_node, + self.inferred_signature(db), + self.inferred_defaults(db), has_implicitly_positional_first_parameter, ); diff --git a/crates/ty_python_semantic/src/types/infer.rs b/crates/ty_python_semantic/src/types/infer.rs index a68182eabf..ccc0ceae1a 100644 --- a/crates/ty_python_semantic/src/types/infer.rs +++ b/crates/ty_python_semantic/src/types/infer.rs @@ -49,9 +49,8 @@ use crate::semantic_index::expression::Expression; use crate::semantic_index::scope::ScopeId; use crate::semantic_index::{SemanticIndex, semantic_index}; use crate::types::diagnostic::TypeCheckDiagnostics; -use crate::types::function::FunctionType; +use crate::types::function::{FunctionType, InferredFunctionSignature}; use crate::types::generics::Specialization; -use crate::types::signatures::Parameter; use crate::types::unpacker::{UnpackResult, Unpacker}; use crate::types::{ ClassLiteral, KnownClass, Truthiness, Type, TypeAndQualifiers, declaration_type, @@ -527,12 +526,6 @@ impl<'db> InferenceRegion<'db> { } } -#[derive(Debug, Eq, PartialEq, get_size2::GetSize, salsa::Update)] -struct InferredFunctionSignature<'db> { - parameters: Vec>, - return_type: Option>, -} - /// The inferred types for a scope region. #[derive(Debug, Eq, PartialEq, salsa::Update, get_size2::GetSize)] pub(crate) struct ScopeInference<'db> { @@ -616,13 +609,10 @@ impl<'db> ScopeInference<'db> { extra.string_annotations.contains(&expression.into()) } - pub(crate) fn inferred_function_signature( - &self, - ) -> Option<(Vec>, Option>)> { + pub(crate) fn inferred_function_signature(&self) -> Option> { self.extra .as_ref() - .and_then(|extra| extra.signature.as_ref()) - .map(|signature| (signature.parameters.clone(), signature.return_type)) + .and_then(|extra| extra.signature.as_ref().cloned()) } } @@ -801,13 +791,10 @@ impl<'db> DefinitionInference<'db> { self.extra.as_ref().and_then(|extra| extra.undecorated_type) } - pub(crate) fn inferred_function_signature( - &self, - ) -> Option<(Vec>, Option>)> { + pub(crate) fn inferred_function_signature(&self) -> Option> { self.extra .as_ref() - .and_then(|extra| extra.signature.as_ref()) - .map(|signature| (signature.parameters.clone(), signature.return_type)) + .and_then(|extra| extra.signature.as_ref().cloned()) } } diff --git a/crates/ty_python_semantic/src/types/infer/builder.rs b/crates/ty_python_semantic/src/types/infer/builder.rs index e8ab547c87..db5dcdde89 100644 --- a/crates/ty_python_semantic/src/types/infer/builder.rs +++ b/crates/ty_python_semantic/src/types/infer/builder.rs @@ -16,9 +16,9 @@ use smallvec::SmallVec; use super::{ DefinitionInference, DefinitionInferenceExtra, ExpressionInference, ExpressionInferenceExtra, - InferenceRegion, InferredFunctionSignature, ScopeInference, ScopeInferenceExtra, - infer_deferred_types, infer_definition_types, infer_expression_types, - infer_same_file_expression_type, infer_unpack_types, + InferenceRegion, ScopeInference, ScopeInferenceExtra, infer_deferred_types, + infer_definition_types, infer_expression_types, infer_same_file_expression_type, + infer_unpack_types, }; use crate::diagnostic::format_enumeration; use crate::module_name::{ModuleName, ModuleNameResolutionError}; @@ -83,8 +83,8 @@ use crate::types::diagnostic::{ report_rebound_typevar, report_slice_step_size_zero, report_unsupported_comparison, }; use crate::types::function::{ - FunctionDecorators, FunctionLiteral, FunctionType, KnownFunction, OverloadLiteral, - is_implicit_classmethod, is_implicit_staticmethod, + FunctionDecorators, FunctionLiteral, FunctionType, InferredFunctionSignature, KnownFunction, + OverloadLiteral, is_implicit_classmethod, is_implicit_staticmethod, }; use crate::types::generics::{ GenericContext, InferableTypeVars, LegacyGenericBase, SpecializationBuilder, bind_typevar, @@ -2289,7 +2289,7 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { decorator_types_and_nodes.push((decorator_type, decorator)); } - let _defaults: Vec<_> = parameters + let defaults: Vec<_> = parameters .iter_non_variadic_params() .map(|param| { param @@ -2301,9 +2301,10 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { // If there are type params, parameters and returns are evaluated in that scope, that is, in // `infer_function_type_params`, rather than here. - if type_params.is_none() { + let inferred_signature = if type_params.is_none() { if self.defer_annotations() { self.deferred.insert(definition, self.multi_inference_state); + None } else { let previous_typevar_binding_context = self.typevar_binding_context.replace(definition); @@ -2317,8 +2318,11 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { return_type, }); self.typevar_binding_context = previous_typevar_binding_context; + self.signature.clone() } - } + } else { + None + }; let known_function = KnownFunction::try_from_definition_and_name(self.db(), definition, name); @@ -2341,6 +2345,8 @@ impl<'db, 'ast> TypeInferenceBuilder<'db, 'ast> { function_decorators, deprecated, dataclass_transformer_params, + inferred_signature, + defaults, ); let function_literal = FunctionLiteral::new(self.db(), overload_literal); diff --git a/crates/ty_python_semantic/src/types/signatures.rs b/crates/ty_python_semantic/src/types/signatures.rs index 916511c5b8..6f817715ac 100644 --- a/crates/ty_python_semantic/src/types/signatures.rs +++ b/crates/ty_python_semantic/src/types/signatures.rs @@ -15,14 +15,18 @@ use std::{collections::HashMap, slice::Iter}; use itertools::EitherOrBoth; use smallvec::{SmallVec, smallvec_inline}; -use super::{DynamicType, Type, TypeVarVariance, definition_expression_type}; +use super::{DynamicType, Type, TypeVarVariance}; use crate::semantic_index::definition::Definition; +use crate::semantic_index::scope::NodeWithScopeRef; +use crate::semantic_index::semantic_index; use crate::types::constraints::{ConstraintSet, IteratorConstraintsExtension}; +use crate::types::function::InferredFunctionSignature; use crate::types::generics::{GenericContext, InferableTypeVars, walk_generic_context}; use crate::types::{ ApplyTypeMappingVisitor, BindingContext, BoundTypeVarInstance, FindLegacyTypeVarsVisitor, HasRelationToVisitor, IsDisjointVisitor, IsEquivalentVisitor, KnownClass, MaterializationKind, - NormalizedVisitor, TypeContext, TypeMapping, TypeRelation, VarianceInferable, todo_type, + NormalizedVisitor, TypeContext, TypeMapping, TypeRelation, VarianceInferable, + infer_deferred_types, infer_scope_types, todo_type, }; use crate::{Db, FxOrderSet}; use ruff_python_ast::{self as ast, name::Name}; @@ -395,18 +399,41 @@ impl<'db> Signature<'db> { pep695_generic_context: Option>, definition: Definition<'db>, function_node: &ast::StmtFunctionDef, + inferred_signature: Option>, + defaults: &[Option>], has_implicitly_positional_first_parameter: bool, ) -> Self { - let mut parameters = - Parameters::from_parameters(db, definition, function_node.parameters.as_ref()); + let mut inferred_signature = inferred_signature + .or_else(|| { + let file = definition.file(db); + let index = semantic_index(db, file); + let file_scope = index + .try_node_scope(NodeWithScopeRef::FunctionTypeParameters(function_node))?; + let scope = file_scope.to_scope_id(db, file); + infer_scope_types(db, scope).inferred_function_signature() + }) + .or_else(|| infer_deferred_types(db, definition).inferred_function_signature()) + .expect("should have an inferred function signature"); + + let mut defaults = defaults.into_iter(); + for parameter in &mut inferred_signature.parameters { + parameter.update_default_type(|| { + defaults + .next() + .expect("should have optional default for each non-variadic parameter") + .map(|ty| ty.replace_parameter_defaults(db)) + }); + } + let mut parameters = Parameters::new(db, inferred_signature.parameters); parameters .find_pep484_positional_only_parameters(has_implicitly_positional_first_parameter); - let return_ty = function_node - .returns - .as_ref() - .map(|returns| definition_expression_type(db, definition, returns.as_ref())); - let legacy_generic_context = - GenericContext::from_function_params(db, definition, ¶meters, return_ty); + + let legacy_generic_context = GenericContext::from_function_params( + db, + definition, + ¶meters, + inferred_signature.return_type, + ); let full_generic_context = GenericContext::merge_pep695_and_legacy( db, pep695_generic_context, @@ -417,7 +444,7 @@ impl<'db> Signature<'db> { generic_context: full_generic_context, definition: Some(definition), parameters, - return_ty, + return_ty: inferred_signature.return_type, } } @@ -1394,96 +1421,6 @@ impl<'db> Parameters<'db> { } } - fn from_parameters( - db: &'db dyn Db, - definition: Definition<'db>, - parameters: &ast::Parameters, - ) -> Self { - let ast::Parameters { - posonlyargs, - args, - vararg, - kwonlyargs, - kwarg, - range: _, - node_index: _, - } = parameters; - - let default_type = |param: &ast::ParameterWithDefault| { - param.default().map(|default| { - definition_expression_type(db, definition, default).replace_parameter_defaults(db) - }) - }; - - let positional_only = posonlyargs.iter().map(|arg| { - Parameter::from_node_and_kind( - db, - definition, - &arg.parameter, - ParameterKind::PositionalOnly { - name: Some(arg.parameter.name.id.clone()), - default_type: default_type(arg), - }, - ) - }); - - let positional_or_keyword = args.iter().map(|arg| { - Parameter::from_node_and_kind( - db, - definition, - &arg.parameter, - ParameterKind::PositionalOrKeyword { - name: arg.parameter.name.id.clone(), - default_type: default_type(arg), - }, - ) - }); - - let variadic = vararg.as_ref().map(|arg| { - Parameter::from_node_and_kind( - db, - definition, - arg, - ParameterKind::Variadic { - name: arg.name.id.clone(), - }, - ) - }); - - let keyword_only = kwonlyargs.iter().map(|arg| { - Parameter::from_node_and_kind( - db, - definition, - &arg.parameter, - ParameterKind::KeywordOnly { - name: arg.parameter.name.id.clone(), - default_type: default_type(arg), - }, - ) - }); - - let keywords = kwarg.as_ref().map(|arg| { - Parameter::from_node_and_kind( - db, - definition, - arg, - ParameterKind::KeywordVariadic { - name: arg.name.id.clone(), - }, - ) - }); - - Self::new( - db, - positional_only - .into_iter() - .chain(positional_or_keyword) - .chain(variadic) - .chain(keyword_only) - .chain(keywords), - ) - } - pub(crate) fn find_pep484_positional_only_parameters( &mut self, has_implicitly_positional_first_parameter: bool, @@ -1724,6 +1661,15 @@ impl<'db> Parameter<'db> { self } + pub(crate) fn update_default_type(&mut self, default: impl FnOnce() -> Option>) { + match &mut self.kind { + ParameterKind::PositionalOnly { default_type, .. } + | ParameterKind::PositionalOrKeyword { default_type, .. } + | ParameterKind::KeywordOnly { default_type, .. } => *default_type = default(), + ParameterKind::Variadic { .. } | ParameterKind::KeywordVariadic { .. } => {} + } + } + pub(crate) fn type_form(mut self) -> Self { self.form = ParameterForm::Type; self @@ -1883,22 +1829,6 @@ impl<'db> Parameter<'db> { }) } - fn from_node_and_kind( - db: &'db dyn Db, - definition: Definition<'db>, - parameter: &ast::Parameter, - kind: ParameterKind<'db>, - ) -> Self { - Self { - annotated_type: parameter - .annotation() - .map(|annotation| definition_expression_type(db, definition, annotation)), - kind, - form: ParameterForm::Value, - inferred_annotation: false, - } - } - /// Returns `true` if this is a keyword-only parameter. pub(crate) fn is_keyword_only(&self) -> bool { matches!(self.kind, ParameterKind::KeywordOnly { .. })