From 64950aa37fb25f4efa5c7606e52b70059c080db3 Mon Sep 17 00:00:00 2001 From: Kevin Date: Sat, 14 Nov 2026 12:48:55 +0800 Subject: [PATCH] Narrow TypeGuard/TypeIs subjects bound by keyword argument --- pyrefly/lib/alt/call.rs | 97 +++++- pyrefly/lib/alt/callable.rs | 196 ++++++++---- pyrefly/lib/alt/class/dataclass.rs | 1 + pyrefly/lib/alt/narrow.rs | 473 ++++++++++++++++------------- pyrefly/lib/alt/overload.rs | 75 +++-- pyrefly/lib/binding/narrow.rs | 267 ++++++++++++---- pyrefly/lib/test/narrow.rs | 188 ++++++++++++ 7 files changed, 944 insertions(+), 353 deletions(-) diff --git a/pyrefly/lib/alt/call.rs b/pyrefly/lib/alt/call.rs index ea77018fb5..a4fad79180 100644 --- a/pyrefly/lib/alt/call.rs +++ b/pyrefly/lib/alt/call.rs @@ -61,6 +61,7 @@ use crate::alt::unwrap::HintRef; use crate::alt::unwrap::MAX_HINT_WIDTH; use crate::binding::binding::Key; use crate::binding::binding::LambdaParamId; +use crate::binding::narrow::GuardSubject; use crate::config::error_kind::ErrorKind; use crate::error::collector::ErrorCollector; use crate::error::context::ErrorContext; @@ -173,6 +174,10 @@ pub struct IntersectionCallTarget { pub struct CallOutcome { pub ty: Type, pub overload_table: OverloadTable, + /// The argument supplying the designated parameter of a `TypeGuard`/`TypeIs` result, + /// identified by the same call resolution that produced `ty`; `None` when the call + /// has no guard result or no single argument could be identified. + pub guard_subject: Option, } impl CallOutcome { @@ -180,6 +185,7 @@ impl CallOutcome { Self { ty, overload_table: OverloadTable::default(), + guard_subject: None, } } } @@ -1493,6 +1499,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { let CallOutcome { ty: ret, overload_table, + .. } = self.call_infer_with_callee_range( self.dunder_new_call_target( new_method.clone(), @@ -2045,6 +2052,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { ) => self.call_infer_inner( signature, Some(&metadata.kind), + 1, tparams.as_deref(), metadata.flags.unresolved_imports, Some(obj), @@ -2061,6 +2069,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { CallTarget::Callable(TargetWithTParams(tparams, callable)) => self.call_infer_inner( callable, None, + 0, tparams.as_deref(), None, None, @@ -2083,6 +2092,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { )) => self.call_infer_inner( callable, Some(&metadata.kind), + Self::receiver_param_count(&metadata), tparams.as_deref(), metadata.flags.unresolved_imports, None, @@ -2097,9 +2107,10 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { ctor_targs, ), CallTarget::FunctionOverload(overloads, metadata) => { - let (ty, _callable, overload_table) = self.call_overloads( + let (ty, _callable, overload_table, guard_subject) = self.call_overloads( overloads, &metadata, + Self::receiver_param_count(&metadata), None, args, keywords, @@ -2111,12 +2122,17 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { hint, ctor_targs, ); - CallOutcome { ty, overload_table } + CallOutcome { + ty, + overload_table, + guard_subject, + } } CallTarget::BoundMethodOverload(obj, overloads, meta) => { - let (ty, _callable, overload_table) = self.call_overloads( + let (ty, _callable, overload_table, guard_subject) = self.call_overloads( overloads, &meta, + 1, Some(obj), args, keywords, @@ -2128,13 +2144,17 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { hint, ctor_targs, ); - CallOutcome { ty, overload_table } + CallOutcome { + ty, + overload_table, + guard_subject, + } } CallTarget::Union(targets) => { let call = CallWithTypes::new(); let args = call.vec_call_arg(args, self, errors); let keywords = call.vec_call_keyword(keywords, self, errors); - let ty = self.unions(targets.into_map(|t| { + let outcomes = targets.into_map(|t| { let ctor_targs = None; // hack self.call_infer_with_callee_range( t, @@ -2149,9 +2169,29 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { hint, ctor_targs, ) - .ty - })); - CallOutcome::of_ty(ty) + }); + let ty = self.unions( + outcomes + .iter() + .map(|outcome| outcome.ty.clone()) + .collect::>(), + ); + // Applicable alternatives that agree on the guarded argument keep it; when + // they disagree, no guarded argument is identified and nothing narrows. + let first = outcomes[0].guard_subject.clone(); + let guard_subject = if outcomes + .iter() + .all(|outcome| outcome.guard_subject == first) + { + first + } else { + None + }; + CallOutcome { + ty, + overload_table: OverloadTable::default(), + guard_subject, + } } CallTarget::Intersect(targets) => self.call_intersection( targets, @@ -2184,6 +2224,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { let CallOutcome { ty: res, overload_table, + guard_subject, } = outcome; let res = if let Some((func_metadata, wrapped)) = wrapped_arg.as_ref() && func_metadata.kind == FunctionKind::FunctoolsWraps @@ -2251,9 +2292,26 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { CallOutcome { ty: res, overload_table, + guard_subject, } } + /// The number of leading receiver parameters of the function, when it is known to be + /// a method: the receiver occupies `params[0]` whether it is bound by attribute access + /// or supplied explicitly. Returns 0 for non-methods and, together with the caller's + /// guard-subject computation, declines to narrow when the information is unavailable. + fn receiver_param_count(metadata: &FuncMetadata) -> usize { + if metadata.flags.is_staticmethod { + return 0; + } + let defined_in_class = match &metadata.kind { + FunctionKind::Def(id) => id.cls.is_some(), + FunctionKind::Synthesized(symbol) => symbol.cls.is_some(), + _ => false, + }; + usize::from(defined_in_class) + } + fn constructor_witness_matches(&self, witness: &Type, class: &ClassType) -> bool { match witness { Type::ClassType(bound) | Type::SelfType(bound) => { @@ -2591,9 +2649,10 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { .snapshot_reachable_vars(&roots.iter().collect::>()); let lambda_snapshot = self.snapshot_lambda_param_types(); let errors = self.error_collector(); - let (_, specialization_errors, _, probe_args, _, _) = self.callable_infer( + let (_, specialization_errors, _, probe_args, _, _, _) = self.callable_infer( probe, None, + 0, tparams, self_obj.cloned(), args, @@ -2664,6 +2723,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { &self, callable: Callable, callable_name: Option<&FunctionKind>, + receiver_offset: usize, tparams: Option<&TParams>, declared_unresolved_imports: Option, self_obj: Option, @@ -2700,6 +2760,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { let res_no_hint = self.callable_infer( callable, callable_name, + receiver_offset, tparams, self_obj, args, @@ -2723,6 +2784,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { let res_with_hint = self.callable_infer( callable, callable_name, + receiver_offset, tparams, self_obj, args, @@ -2783,8 +2845,15 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { { *targs = chosen_targs; } - let (ty, specialization_errors, return_type_errors, _expected_types, _, overload_table) = - chosen_res; + let ( + ty, + specialization_errors, + return_type_errors, + _expected_types, + _, + overload_table, + guard_subject, + ) = chosen_res; if let Ok(errors) = Vec1::try_from_vec(specialization_errors) { self.add_specialization_errors(errors, arguments_range, call_errors, context); } @@ -2798,7 +2867,11 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { // Keep probe failures out of the real call's contextual typing and operator dispatch. self.warn_unresolved_import_signature(return_errors, arguments_range, input_errors); } - CallOutcome { ty, overload_table } + CallOutcome { + ty, + overload_table, + guard_subject, + } } pub fn call_infer( diff --git a/pyrefly/lib/alt/callable.rs b/pyrefly/lib/alt/callable.rs index cdb5089059..e9c2eda43f 100644 --- a/pyrefly/lib/alt/callable.rs +++ b/pyrefly/lib/alt/callable.rs @@ -54,6 +54,7 @@ use crate::alt::shape_extension::shape_extension_vars; use crate::alt::solve::Iterable; use crate::alt::unwrap::HintRef; use crate::alt::unwrap::MAX_HINT_WIDTH; +use crate::binding::narrow::GuardSubject; use crate::config::error_kind::ErrorKind; use crate::error::collector::ErrorCollector; use crate::error::context::ErrorContext; @@ -942,6 +943,51 @@ impl ArgMap { } } +/// Identify which argument supplies a `TypeGuard`/`TypeIs` function's designated +/// parameter: the first non-receiver positional-or-keyword parameter. The +/// argument-to-parameter mapping comes from `callable_infer_params`, so keyword binding +/// and keyword order follow the standard call logic; `receiver_offset` skips the receiver +/// parameter of a method independently of how the receiver was supplied. +/// +/// Returns `None` when no single argument can be identified: the signature is gradual, +/// the designated parameter is missing or not keyword-passable, or it was supplied by a +/// default or a `**kwargs` splat rather than one definite argument. +fn guarded_param_subject( + params: &ParamList, + receiver_offset: usize, + has_self_arg: bool, + argmap: &ArgMap, + keywords: &[CallKeyword], +) -> Option { + let designated = params + .items() + .get(receiver_offset) + .and_then(|param| match param { + Param::Pos(name, _, _) | Param::PosOnly(Some(name), _, _) => Some(name), + // Positional-only parameters without a name can't be looked up, and a + // `TypeGuard`/`TypeIs` designation never lands on other parameter kinds. + _ => None, + })?; + let (argument, _) = argmap + .arg_to_param + .iter() + .find(|(_, matched)| matched.name.as_ref() == Some(designated))?; + match argument { + ArgumentKey::Positional(index) => { + // The receiver argument itself is never the designated parameter. + (*index) + .checked_sub(usize::from(has_self_arg)) + .map(GuardSubject::Positional) + } + ArgumentKey::Keyword(index) => { + let kw = keywords.get(*index)?; + kw.arg + .filter(|id| &id.id == designated) + .map(|id| GuardSubject::Keyword(id.id.clone())) + } + } +} + /// Helps track matching of arguments against positional parameters in AnswersSolver::callable_infer_params. #[derive(PartialEq, Eq)] enum PosParamKind { @@ -2473,6 +2519,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { &self, callable: Callable, callable_name: Option<&FunctionKind>, + receiver_offset: usize, tparams: Option<&TParams>, self_obj: Option, args: &[CallArg], @@ -2491,6 +2538,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { ArgMap, SmallSet, OverloadTable, + Option, ) { let hint = HintRef::filter_for_call(hint, tparams); self.callable_infer_with_hint( @@ -2502,6 +2550,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { self.callable_infer_inner( callable, callable_name, + receiver_offset, tparams, self_obj, args, @@ -2523,6 +2572,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { &self, callable: Callable, callable_name: Option<&FunctionKind>, + receiver_offset: usize, tparams: Option<&TParams>, mut self_obj: Option, mut args: &[CallArg], @@ -2541,6 +2591,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { ArgMap, SmallSet, OverloadTable, + Option, ) { let call_boundary = CallBoundary::new(); let call_context = call_boundary @@ -2649,21 +2700,32 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { arguments_range, }); let self_arg = self_obj.as_ref().map(|ty| CallArg::ty(ty, arguments_range)); - let argmap = match callable.params { - Params::List(params) | Params::Partial(params) => self.callable_infer_params( - callable_name, - ¶ms, - None, - self_arg, - &mut self_qs, - args, - keywords, - arguments_range, - arg_errors, - call_errors, - context, - &call_context, - ), + let has_self_arg = self_arg.is_some(); + let (argmap, guard_subject) = match callable.params { + Params::List(params) | Params::Partial(params) => { + let argmap = self.callable_infer_params( + callable_name, + ¶ms, + None, + self_arg, + &mut self_qs, + args, + keywords, + arguments_range, + arg_errors, + call_errors, + context, + &call_context, + ); + let guard_subject = guarded_param_subject( + ¶ms, + receiver_offset, + has_self_arg, + &argmap, + keywords, + ); + (argmap, guard_subject) + } Params::Ellipsis | Params::Materialization => { // Deal with Callable[..., R] for arg in self_arg.iter().chain(args.iter()) { @@ -2672,48 +2734,70 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { for keyword in keywords { keyword.value.infer(self, arg_errors); } - ArgMap::new() + (ArgMap::new(), None) } Params::ParamSpec(concatenate, p) => { let p = self.solver().expand(p); match p { - Type::ParamSpecValue(params) => self.callable_infer_params( - callable_name, - ¶ms.prepend_types(&concatenate), - None, - self_arg, - &mut self_qs, - args, - keywords, - arguments_range, - arg_errors, - call_errors, - context, - &call_context, - ), + Type::ParamSpecValue(params) => { + let matched = params.prepend_types(&concatenate); + let argmap = self.callable_infer_params( + callable_name, + &matched, + None, + self_arg, + &mut self_qs, + args, + keywords, + arguments_range, + arg_errors, + call_errors, + context, + &call_context, + ); + let guard_subject = guarded_param_subject( + &matched, + receiver_offset, + has_self_arg, + &argmap, + keywords, + ); + (argmap, guard_subject) + } // This can happen with a signature like `(f: Callable[P, None], *args: P.args, **kwargs: P.kwargs)`. // Before we match an argument to `f`, we don't know what `P` is, so we don't have an answer for the Var yet. // Use to_subset_param to preserve Pos vs PosOnly: prefix params from a // function definition should remain keyword-passable in direct calls. - Type::Var(var) => self.callable_infer_params( - callable_name, - &ParamList::new( + Type::Var(var) => { + let matched = ParamList::new( concatenate .iter() .map(|p| p.to_param_preserve_name()) .collect(), - ), - Some(var), - self_arg, - &mut self_qs, - args, - keywords, - arguments_range, - arg_errors, - call_errors, - context, - &call_context, - ), + ); + let argmap = self.callable_infer_params( + callable_name, + &matched, + Some(var), + self_arg, + &mut self_qs, + args, + keywords, + arguments_range, + arg_errors, + call_errors, + context, + &call_context, + ); + let guard_subject = guarded_param_subject( + &matched, + receiver_offset, + has_self_arg, + &argmap, + keywords, + ); + (argmap, guard_subject) + } Type::Quantified(q) => { if let Some((args, keywords)) = self.paramspec_forwarding( &q, @@ -2725,9 +2809,10 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { call_errors, context, ) { - self.callable_infer_params( + let matched = ParamList::new_types(concatenate.into_vec()); + let argmap = self.callable_infer_params( callable_name, - &ParamList::new_types(concatenate.into_vec()), + &matched, None, self_arg, &mut self_qs, @@ -2738,12 +2823,20 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { call_errors, context, &call_context, - ) + ); + let guard_subject = guarded_param_subject( + &matched, + receiver_offset, + has_self_arg, + &argmap, + keywords, + ); + (argmap, guard_subject) } else { - ArgMap::new() + (ArgMap::new(), None) } } - Type::Any(_) | Type::Ellipsis => ArgMap::new(), + Type::Any(_) | Type::Ellipsis => (ArgMap::new(), None), _ => { // This could well be our error, but not really sure self.error_with_context( @@ -2753,7 +2846,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { format!("Unexpected ParamSpec type: `{}`", self.for_display(p)), context, ); - ArgMap::new() + (ArgMap::new(), None) } } } @@ -2792,6 +2885,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { argmap, defaults_used, overload_table, + guard_subject, ) } } diff --git a/pyrefly/lib/alt/class/dataclass.rs b/pyrefly/lib/alt/class/dataclass.rs index d4bbc1e50f..96e0c68a4a 100644 --- a/pyrefly/lib/alt/class/dataclass.rs +++ b/pyrefly/lib/alt/class/dataclass.rs @@ -992,6 +992,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { })) .unwrap(), &overload.metadata, + 0, None, &args.args.map(CallArg::expr_maybe_starred), &args.keywords.map(CallKeyword::new), diff --git a/pyrefly/lib/alt/narrow.rs b/pyrefly/lib/alt/narrow.rs index 324730f6c4..66bb8b0c3e 100644 --- a/pyrefly/lib/alt/narrow.rs +++ b/pyrefly/lib/alt/narrow.rs @@ -26,7 +26,6 @@ use pyrefly_types::type_info::JoinStyle; use pyrefly_types::typed_dict::ExtraItems; use pyrefly_util::prelude::SliceExt; use pyrefly_util::visit::Visit; -use ruff_python_ast::Arguments; use ruff_python_ast::AtomicNodeIndex; use ruff_python_ast::Expr; use ruff_python_ast::ExprBinOp; @@ -54,10 +53,14 @@ use crate::alt::types::instance::Instance; use crate::binding::binding::Key; use crate::binding::binding::MatchCaseNegations; use crate::binding::narrow::AtomicNarrowOp; +use crate::binding::narrow::CallCandidate; use crate::binding::narrow::FacetOrigin; use crate::binding::narrow::FacetSubject; +use crate::binding::narrow::GuardNarrowing; +use crate::binding::narrow::GuardSubject; use crate::binding::narrow::NarrowOp; use crate::binding::narrow::NarrowSource; +use crate::binding::narrow::NarrowingCall; use crate::binding::narrow::NarrowingSubject; use crate::error::collector::ErrorCollector; use crate::error::style::ErrorStyle; @@ -98,6 +101,16 @@ fn synthesize_int_slice(idx: i64) -> Expr { } /// Append `facet` to an existing facet chain, or start a new one. +/// Whether a recorded candidate argument is the one that standard call resolution +/// identified as supplying the `TypeGuard`/`TypeIs` designated parameter. A positional +/// candidate (`kw == None`) stands for the call's first positional argument. +fn candidate_matches(candidate: &CallCandidate, subject: &GuardSubject) -> bool { + match subject { + GuardSubject::Keyword(name) => candidate.kw.as_ref() == Some(name), + GuardSubject::Positional(index) => candidate.kw.is_none() && *index == 0, + } +} + fn extend_facet_chain(resolved_chain: Option<&FacetChain>, facet: FacetKind) -> Vec1 { match resolved_chain { Some(chain) => { @@ -683,13 +696,13 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { fn resolve_narrowing_call( &self, - func: &Expr, - args: &Arguments, + call: &NarrowingCall, + candidates: &[CallCandidate], errors: &ErrorCollector, ) -> Option { - let func_ty = self.expr_infer(func, errors); - if args.args.len() > 1 { - let second_arg = &args.args[1]; + let func_ty = self.expr_infer(&call.func, errors); + if call.args.args.len() > 1 { + let second_arg = &call.args.args[1]; let op = match func_ty.callee_kind() { Some(CalleeKind::Function(FunctionKind::IsInstance)) => Some( AtomicNarrowOp::IsInstance(second_arg.clone(), NarrowSource::Call), @@ -703,12 +716,67 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { return op; } } - if func_ty.is_typeis() { - Some(AtomicNarrowOp::TypeIs(func_ty.clone(), args.clone())) - } else if func_ty.is_typeguard() { - Some(AtomicNarrowOp::TypeGuard(func_ty.clone(), args.clone())) - } else { - None + // Cheaply reject callees that cannot produce a `TypeGuard`/`TypeIs` result before + // paying for full call resolution. + if !func_ty.is_typeis() && !func_ty.is_typeguard() { + return None; + } + let CallTargetLookup::Ok(call_target) = self.as_call_target(func_ty.clone()) else { + return None; + }; + let args = call.args.args.map(CallArg::expr_maybe_starred); + let keywords = call.args.keywords.map(CallKeyword::new); + // The resolution errors are recorded in this throwaway collector, and dropped with + // it because they are reported by the own inference of the call. Recording them + // lets a call that does not go through cleanly be detected here. + let thrown_away_errors = self.error_collector(); + let outcome = self.call_infer( + *call_target, + &args, + &keywords, + call.args.range(), + &thrown_away_errors, + None, + None, + None, + ); + // A keyword-supplied candidate argument only narrows when the call goes through + // cleanly: a call that reports any error (for example a keyword bound to a + // positional-only parameter) raises `TypeError` at runtime and does not narrow. + if !thrown_away_errors.is_empty() && candidates.iter().any(|c| c.kw.is_some()) { + return None; + } + // The narrowed type and the guarded argument come from the same call resolution; a + // call whose resolution lands on a non-guard result (for example through overloads) + // does not narrow. + match outcome.ty { + Type::TypeIs(narrowed) => { + let target = if func_ty.callee_kind() + == Some(CalleeKind::Function(FunctionKind::Callable)) + { + // `callable` is annotated as `TypeIs[Callable[..., object]]` which is + // too conservative and prone to false positives, see + // https://github.com/facebook/pyrefly/issues/911 + self.heap.mk_callable_ellipsis(self.heap.mk_any_implicit()) + } else { + *narrowed + }; + Some(AtomicNarrowOp::TypeIs(Box::new(GuardNarrowing { + narrowed: target, + func_ty, + args: call.args.clone(), + subject: outcome.guard_subject, + }))) + } + Type::TypeGuard(narrowed) => { + Some(AtomicNarrowOp::TypeGuard(Box::new(GuardNarrowing { + narrowed: *narrowed, + func_ty, + args: call.args.clone(), + subject: outcome.guard_subject, + }))) + } + _ => None, } } @@ -1618,6 +1686,112 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { } } + /// Apply one facet-aware narrow: resolve the facet chain, narrow the facet type, + /// record it on the subject, and narrow the subject itself through the facet when + /// applicable (a facet narrow of the last facet element also filters the subject). + fn narrow_by_facet( + &self, + type_info: &TypeInfo, + facet_subject: &FacetSubject, + op: &AtomicNarrowOp, + range: TextRange, + errors: &ErrorCollector, + ) -> TypeInfo { + let Some(resolved_chain) = self.resolve_facet_chain(facet_subject.chain.clone()) else { + return type_info.clone(); + }; + match (facet_subject.origin, resolved_chain.facets().as_slice()) { + (FacetOrigin::GetMethod, _) + if !self.supports_dict_get_subject(type_info, facet_subject, range) => + { + return type_info.clone(); + } + (FacetOrigin::MatchSubject, [FacetKind::Index(index)]) => { + let index = usize::try_from(*index) + .expect("Match subject indices are nonnegative tuple positions"); + // Keep element constraints in the evaluated tuple's type so that + // joining alternatives preserves their correlation across cases. + // For example, excluding (None, None) leaves a union of tuples + // with either the first or the second element known to be present. + let ty = self.distribute_over_union(type_info.ty(), |ty| match ty { + Type::Tuple(Tuple::Concrete(elements)) if index < elements.len() => { + match self.atomic_narrow(&elements[index], op, range, errors) { + narrowed @ Type::Never(_) => narrowed, + narrowed => { + let mut elements = elements.clone(); + elements[index] = narrowed; + self.heap.mk_concrete_tuple(elements) + } + } + } + _ => ty.clone(), + }); + return type_info.clone().with_ty(ty); + } + _ => {} + } + let ty = self.atomic_narrow( + &self.get_facet_chain_type(type_info, &resolved_chain, range), + op, + range, + errors, + ); + // A match-arm negation may build on a facet narrow from an earlier arm. + // If the accumulated facet is impossible, the whole subject is impossible. + if facet_subject.origin == FacetOrigin::MatchSubject + && facet_subject.allow_never_collapse + && ty.is_never() + { + return type_info.clone().with_ty(ty); + } + let mut narrowed = type_info.with_narrow(resolved_chain.facets(), ty); + // For certain types of narrows, we can also narrow the parent of the current subject + // If `.get()` on a dict or TypedDict is falsy, the key may not be present at all + // We should invalidate any existing narrows + if let Some((last, prefix)) = resolved_chain.facets().split_last() { + match Vec1::try_from(prefix) { + Ok(prefix_facets) => { + let prefix_chain = FacetChain::new(prefix_facets); + let base_ty = self.get_facet_chain_type(type_info, &prefix_chain, range); + let dict_get_key_falsy = + matches!(op, AtomicNarrowOp::IsFalsy) && matches!(last, FacetKind::Key(_)); + if dict_get_key_falsy { + narrowed.update_for_assignment(resolved_chain.facets(), None); + } else if let Some(narrowed_ty) = self.atomic_narrow_for_facet( + &base_ty, + last, + op, + facet_subject.allow_never_collapse, + range, + errors, + ) && narrowed_ty != base_ty + { + narrowed = narrowed.with_narrow(prefix_chain.facets(), narrowed_ty); + } + } + _ => { + let base_ty = type_info.ty(); + let dict_get_key_falsy = + matches!(op, AtomicNarrowOp::IsFalsy) && matches!(last, FacetKind::Key(_)); + if dict_get_key_falsy { + narrowed.update_for_assignment(resolved_chain.facets(), None); + } else if let Some(narrowed_ty) = self.atomic_narrow_for_facet( + base_ty, + last, + op, + facet_subject.allow_never_collapse, + range, + errors, + ) && narrowed_ty != *base_ty + { + narrowed = narrowed.clone().with_ty(narrowed_ty); + } + } + }; + } + narrowed + } + pub(crate) fn atomic_narrow( &self, ty: &Type, @@ -1627,6 +1801,10 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { ) -> Type { match op { AtomicNarrowOp::Placeholder => ty.clone(), + // Call narrowing is applied in `narrow`, which selects the guarded argument + // before producing any facet effect; `atomic_narrow` is only reached for the + // instantiated `TypeIs`/`TypeGuard` forms. + AtomicNarrowOp::Call(..) | AtomicNarrowOp::NotCall(..) => ty.clone(), AtomicNarrowOp::MatchSubjectElement(subject, index) => { let element = self.get_facet_chain_type( self.get_idx(*subject), @@ -1995,89 +2173,14 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { AtomicNarrowOp::NotHasKey(_) => ty.clone(), AtomicNarrowOp::GetAttr(_, _) => ty.clone(), AtomicNarrowOp::NotGetAttr(_, _) => ty.clone(), - AtomicNarrowOp::TypeGuard(t, arguments) => { - if let CallTargetLookup::Ok(call_target) = self.as_call_target(t.clone()) { - let args = arguments.args.map(CallArg::expr_maybe_starred); - let kws = arguments.keywords.map(CallKeyword::new); - // This error is raised elsewhere, swallow here to avoid duplicate errors - let swallowed_errors = self.error_swallower(); - let ret = self - .call_infer( - *call_target, - &args, - &kws, - range, - &swallowed_errors, - None, - None, - None, - ) - .ty; - if let Type::TypeGuard(t) = ret { - return *t; - } - } - ty.clone() - } - AtomicNarrowOp::NotTypeGuard(_, _) => ty.clone(), - AtomicNarrowOp::TypeIs(t, arguments) => { - let is_builtin_callable = - t.callee_kind() == Some(CalleeKind::Function(FunctionKind::Callable)); - if let CallTargetLookup::Ok(call_target) = self.as_call_target(t.clone()) { - let args = arguments.args.map(CallArg::expr_maybe_starred); - let kws = arguments.keywords.map(CallKeyword::new); - // This error is raised elsewhere, swallow here to avoid duplicate errors - let swallowed_errors = self.error_swallower(); - let ret = self - .call_infer( - *call_target, - &args, - &kws, - range, - &swallowed_errors, - None, - None, - None, - ) - .ty; - if let Type::TypeIs(t) = ret { - let target = if is_builtin_callable { - // `callable` is annotated as `TypeIs[Callable[..., object]]` which is - // too conservative and prone to false positives, see - // https://github.com/facebook/pyrefly/issues/911 - self.heap.mk_callable_ellipsis(self.heap.mk_any_implicit()) - } else { - *t - }; - return self.narrow_typeis(ty, &target, range, errors); - } - } - ty.clone() - } - AtomicNarrowOp::NotTypeIs(t, arguments) => { - if let CallTargetLookup::Ok(call_target) = self.as_call_target(t.clone()) { - let args = arguments.args.map(CallArg::expr_maybe_starred); - let kws = arguments.keywords.map(CallKeyword::new); - // This error is raised elsewhere, swallow here to avoid duplicate errors - let swallowed_errors = self.error_swallower(); - let ret = self - .call_infer( - *call_target, - &args, - &kws, - range, - &swallowed_errors, - None, - None, - None, - ) - .ty; - if let Type::TypeIs(t) = ret { - return self.subtract(ty, &t); - } - } - ty.clone() - } + // The call was already resolved when the op was produced; the narrowed type is + // instantiated and carries the identified guarded argument. + AtomicNarrowOp::TypeGuard(narrowing) => narrowing.narrowed.clone(), + AtomicNarrowOp::NotTypeGuard(_) => ty.clone(), + AtomicNarrowOp::TypeIs(narrowing) => { + self.narrow_typeis(ty, &narrowing.narrowed, range, errors) + } + AtomicNarrowOp::NotTypeIs(narrowing) => self.subtract(ty, &narrowing.narrowed), AtomicNarrowOp::IsTruthy | AtomicNarrowOp::IsFalsy => { let boolval = matches!(op, AtomicNarrowOp::IsTruthy); // Do not emit errors here: the narrowed range doesn't always correspond to a valid expression @@ -2166,17 +2269,6 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { ty.clone() } } - AtomicNarrowOp::Call(func, args) | AtomicNarrowOp::NotCall(func, args) => { - if let Some(resolved_op) = self.resolve_narrowing_call(func, args, errors) { - if matches!(op, AtomicNarrowOp::Call(..)) { - self.atomic_narrow(ty, &resolved_op, range, errors) - } else { - self.atomic_narrow(ty, &resolved_op.negate(), range, errors) - } - } else { - ty.clone() - } - } } } @@ -2473,6 +2565,60 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { type_info.with_narrow(&facets, narrowed_ty) } } + NarrowOp::Atomic( + None, + op @ (AtomicNarrowOp::Call(..) | AtomicNarrowOp::NotCall(..)), + ) => { + let (call, candidates, is_call) = match op { + AtomicNarrowOp::Call(call, candidates) => (call, candidates, true), + AtomicNarrowOp::NotCall(call, candidates) => (call, candidates, false), + _ => unreachable!("matched above"), + }; + let Some(resolved_atomic) = self.resolve_narrowing_call(call, candidates, errors) + else { + return type_info.clone(); + }; + // Standard call resolution identified the argument that supplies the + // guarded parameter. Only the recorded candidates matching that argument + // contribute; the other candidates of the same call are absent effects, not + // operands to AND or OR, so both polarities narrow consistently. Calls + // without an identified argument (isinstance-style narrowing of the first + // positional argument) apply to every candidate. + let selected: Vec<&CallCandidate> = match resolved_atomic.guard_subject() { + Some(subject) => candidates + .iter() + .filter(|candidate| candidate_matches(candidate, subject)) + .collect(), + None => candidates.iter().collect(), + }; + if selected.is_empty() { + return type_info.clone(); + } + let mut narrowed_info = type_info.clone(); + for candidate in selected { + let atomic = if is_call { + resolved_atomic.clone() + } else { + resolved_atomic.negate() + }; + match &candidate.facet { + None => { + let ty = self.atomic_narrow( + &self.force_for_narrowing(type_info.ty(), range, errors), + &atomic, + range, + errors, + ); + narrowed_info = self.rebase_type_info(&narrowed_info, ty); + } + Some(facet) => { + narrowed_info = + self.narrow_by_facet(&narrowed_info, facet, &atomic, range, errors); + } + } + } + narrowed_info + } NarrowOp::Atomic(None, op) => { let ty = self.atomic_narrow( &self.force_for_narrowing(type_info.ty(), range, errors), @@ -2483,119 +2629,12 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { self.rebase_type_info(type_info, ty) } NarrowOp::Atomic(Some(facet_subject), op) => { - let Some(resolved_chain) = self.resolve_facet_chain(facet_subject.chain.clone()) - else { - return type_info.clone(); - }; - let Some(op_for_narrow) = (match op { - AtomicNarrowOp::Call(func, args) => { - self.resolve_narrowing_call(func.as_ref(), args, errors) - } - AtomicNarrowOp::NotCall(func, args) => self - .resolve_narrowing_call(func.as_ref(), args, errors) - .map(|resolved_op| resolved_op.negate()), - _ => Some(op.clone()), - }) else { + if matches!(op, AtomicNarrowOp::Call(..) | AtomicNarrowOp::NotCall(..)) { + // Call narrowing is applied in the `None` facet arm above, which + // selects the guarded argument before any facet effect is produced. return type_info.clone(); - }; - match (facet_subject.origin, resolved_chain.facets().as_slice()) { - (FacetOrigin::GetMethod, _) - if !self.supports_dict_get_subject(type_info, facet_subject, range) => - { - return type_info.clone(); - } - (FacetOrigin::MatchSubject, [FacetKind::Index(index)]) => { - let index = usize::try_from(*index) - .expect("Match subject indices are nonnegative tuple positions"); - // Keep element constraints in the evaluated tuple's type so that - // joining alternatives preserves their correlation across cases. - // For example, excluding (None, None) leaves a union of tuples - // with either the first or the second element known to be present. - let ty = self.distribute_over_union(type_info.ty(), |ty| match ty { - Type::Tuple(Tuple::Concrete(elements)) if index < elements.len() => { - match self.atomic_narrow( - &elements[index], - &op_for_narrow, - range, - errors, - ) { - narrowed @ Type::Never(_) => narrowed, - narrowed => { - let mut elements = elements.clone(); - elements[index] = narrowed; - self.heap.mk_concrete_tuple(elements) - } - } - } - _ => ty.clone(), - }); - return type_info.clone().with_ty(ty); - } - _ => {} } - let ty = self.atomic_narrow( - &self.get_facet_chain_type(type_info, &resolved_chain, range), - &op_for_narrow, - range, - errors, - ); - // A match-arm negation may build on a facet narrow from an earlier arm. - // If the accumulated facet is impossible, the whole subject is impossible. - if facet_subject.origin == FacetOrigin::MatchSubject - && facet_subject.allow_never_collapse - && ty.is_never() - { - return type_info.clone().with_ty(ty); - } - let mut narrowed = type_info.with_narrow(resolved_chain.facets(), ty); - // For certain types of narrows, we can also narrow the parent of the current subject - // If `.get()` on a dict or TypedDict is falsy, the key may not be present at all - // We should invalidate any existing narrows - if let Some((last, prefix)) = resolved_chain.facets().split_last() { - match Vec1::try_from(prefix) { - Ok(prefix_facets) => { - let prefix_chain = FacetChain::new(prefix_facets); - let base_ty = - self.get_facet_chain_type(type_info, &prefix_chain, range); - let dict_get_key_falsy = - matches!(op_for_narrow, AtomicNarrowOp::IsFalsy) - && matches!(last, FacetKind::Key(_)); - if dict_get_key_falsy { - narrowed.update_for_assignment(resolved_chain.facets(), None); - } else if let Some(narrowed_ty) = self.atomic_narrow_for_facet( - &base_ty, - last, - &op_for_narrow, - facet_subject.allow_never_collapse, - range, - errors, - ) && narrowed_ty != base_ty - { - narrowed = narrowed.with_narrow(prefix_chain.facets(), narrowed_ty); - } - } - _ => { - let base_ty = type_info.ty(); - let dict_get_key_falsy = - matches!(op_for_narrow, AtomicNarrowOp::IsFalsy) - && matches!(last, FacetKind::Key(_)); - if dict_get_key_falsy { - narrowed.update_for_assignment(resolved_chain.facets(), None); - } else if let Some(narrowed_ty) = self.atomic_narrow_for_facet( - base_ty, - last, - &op_for_narrow, - facet_subject.allow_never_collapse, - range, - errors, - ) && narrowed_ty != *base_ty - { - narrowed = narrowed.clone().with_ty(narrowed_ty); - } - } - }; - } - narrowed + self.narrow_by_facet(type_info, facet_subject, op, range, errors) } NarrowOp::And(ops) => { let mut ops_iter = ops.iter(); diff --git a/pyrefly/lib/alt/overload.rs b/pyrefly/lib/alt/overload.rs index 4fc549a628..e75f552d6e 100644 --- a/pyrefly/lib/alt/overload.rs +++ b/pyrefly/lib/alt/overload.rs @@ -43,6 +43,7 @@ use crate::alt::callable::CallWithTypes; use crate::alt::callable::ReturnTypeResolutionError; use crate::alt::expr::TypeOrExpr; use crate::alt::unwrap::HintRef; +use crate::binding::narrow::GuardSubject; use crate::config::error_kind::ErrorKind; use crate::error::collector::ErrorCollector; use crate::error::context::ErrorContext; @@ -74,6 +75,9 @@ struct CalledOverload<'f> { defaults_used: SmallSet, /// Maps each argument's position to the parameter it was matched against. argmap: ArgMap, + /// The argument supplying this overload's `TypeGuard`/`TypeIs` designated parameter, + /// identified by the same call resolution that produced `res`. + guard_subject: Option, } impl CalledOverload<'_> { @@ -319,6 +323,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { &self, overloads: Vec1>, metadata: &FuncMetadata, + receiver_offset: usize, self_obj: Option, args: &[CallArg], keywords: &[CallKeyword], @@ -330,7 +335,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { hint: Option, // If we're constructing a class, its type arguments. A successful call will fill these in. ctor_targs: Option<&mut TArgs>, - ) -> (Type, Callable, OverloadTable) { + ) -> (Type, Callable, OverloadTable, Option) { // There may be Expr values in args and keywords. // If we infer them for each overload, we may end up inferring them multiple times. // If those overloads contain nested overloads, then we can easily end up with O(2^n) perf. @@ -371,6 +376,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { return_type_errors: Vec::new(), defaults_used: SmallSet::new(), argmap: ArgMap::new(), + guard_subject: None, }, false, ), @@ -380,6 +386,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { let (mut closest_overload, mut matched) = self.find_closest_overload( &arity_compatible_overloads, metadata, + receiver_offset, self_obj.as_ref(), &args, &keywords, @@ -409,6 +416,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { let (cur_closest, cur_matched) = self.find_closest_overload( &arity_compatible_overloads, metadata, + receiver_offset, self_obj.as_ref(), cur_args, cur_keywords, @@ -444,6 +452,16 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { // table of solutions describes the call any more. let table = OverloadTable::default(); let argmap = first_overload.argmap.clone(); + // The alternatives agree on the return, so they may also agree on the + // guarded argument; when they disagree, decline to narrow. + let guard_subject = if matched_overloads + .iter() + .all(|o| o.guard_subject == first_overload.guard_subject) + { + first_overload.guard_subject.clone() + } else { + None + }; let arg_errors = self.error_collector(); let specialization_errors = first_overload.specialization_errors.clone(); let return_type_errors = matched_overloads @@ -469,6 +487,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { specialization_errors, return_type_errors, defaults_used, + guard_subject, }; matched = true; break; @@ -565,6 +584,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { closest_overload.res, closest_overload.func.1.signature.clone(), closest_overload.table, + closest_overload.guard_subject, ) } else { if let Ok(specialization_errors) = @@ -594,6 +614,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { self.heap.mk_any_error(), closest_overload.func.1.signature.clone(), OverloadTable::default(), + None, ) } } @@ -868,6 +889,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { &self, overloads: &Vec1<&'c TargetWithTParams>, metadata: &FuncMetadata, + receiver_offset: usize, self_obj: Option<&Type>, args: &[CallArg], keywords: &[CallKeyword], @@ -887,6 +909,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { let mut called_overload = self.call_overload( callable, metadata, + receiver_offset, self_obj, args, keywords, @@ -1031,6 +1054,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { let res = self.call_overload( o.func, metadata, + receiver_offset, self_obj, &materialized_args, &materialized_keywords, @@ -1063,6 +1087,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { let contextual_overload = self.call_overload( overload.func, metadata, + receiver_offset, self_obj, args, keywords, @@ -1159,6 +1184,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { &self, callable: &'c TargetWithTParams, metadata: &FuncMetadata, + receiver_offset: usize, self_obj: Option<&Type>, args: &[CallArg], keywords: &[CallKeyword], @@ -1176,25 +1202,33 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { let arg_errors = self.error_collector(); let call_errors = self.error_collector(); - let (res, specialization_errors, return_type_errors, argmap, defaults_used, table) = self - .callable_infer( - callable.1.signature.clone(), - Some(&metadata.kind), - tparams, - self_obj.cloned(), - args, - keywords, - arguments_range, - &arg_errors, - &call_errors, - // We intentionally drop the context here, as arg errors don't need it, - // and if there are any call errors, we'll log a "No matching overloads" - // error with the necessary context. - None, - hint, - contextually_opaque_defaults, - overload_ctor_targs.as_mut(), - ); + let ( + res, + specialization_errors, + return_type_errors, + argmap, + defaults_used, + table, + guard_subject, + ) = self.callable_infer( + callable.1.signature.clone(), + Some(&metadata.kind), + receiver_offset, + tparams, + self_obj.cloned(), + args, + keywords, + arguments_range, + &arg_errors, + &call_errors, + // We intentionally drop the context here, as arg errors don't need it, + // and if there are any call errors, we'll log a "No matching overloads" + // error with the necessary context. + None, + hint, + contextually_opaque_defaults, + overload_ctor_targs.as_mut(), + ); if !call_errors.has_hard() && specialization_errors.is_empty() && callable @@ -1226,6 +1260,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { return_type_errors, defaults_used, argmap, + guard_subject, } } } diff --git a/pyrefly/lib/binding/narrow.rs b/pyrefly/lib/binding/narrow.rs index 778b76f8ab..b83df31c44 100644 --- a/pyrefly/lib/binding/narrow.rs +++ b/pyrefly/lib/binding/narrow.rs @@ -6,6 +6,7 @@ */ use std::fmt; +use std::sync::Arc; use pyrefly_graph::index::Idx; use pyrefly_python::ast::Ast; @@ -49,8 +50,8 @@ use crate::types::facet::UnresolvedFacetChain; use crate::types::facet::UnresolvedFacetKind; use crate::types::types::Type; -assert_words!(AtomicNarrowOp, 9); -assert_words!(NarrowOp, 11); +assert_words!(AtomicNarrowOp, 8); +assert_words!(NarrowOp, 10); /// Indicates where an isinstance-style narrow operation originated from. /// This determines whether validation needs to happen during narrowing. @@ -62,6 +63,49 @@ pub enum NarrowSource { Pattern, } +/// A call that is recorded as a possible `TypeGuard`/`TypeIs` guard. +#[derive(Clone, Debug)] +pub struct NarrowingCall { + pub func: Box, + pub args: Arguments, +} + +/// The actual argument that supplies the designated parameter of a `TypeGuard`/`TypeIs` +/// result, as identified by standard call resolution. +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum GuardSubject { + /// A positional argument, indexed from the first argument after the receiver. + Positional(usize), + /// A named keyword argument. + Keyword(Name), +} + +/// A resolved `TypeGuard`/`TypeIs` call: the instantiated guard argument that the subject +/// narrows to, plus the original call for display. +#[derive(Clone, Debug)] +pub struct GuardNarrowing { + /// The instantiated argument of the guard result that the subject narrows to. + pub narrowed: Type, + /// The original callee, for display. + pub func_ty: Type, + pub args: Arguments, + /// The guarded argument identified by standard call resolution, when it identified one. + /// `None` means the narrowing is declined (or, for isinstance-style narrowing, that + /// every candidate applies). + pub subject: Option, +} + +/// One actual argument of a narrowing call that may be the subject its guard result +/// constrains. +#[derive(Clone, Debug)] +pub struct CallCandidate { + /// The keyword supplying the candidate argument, or `None` for the first positional + /// argument. + pub kw: Option, + /// The facet subject of the candidate's value expression, when it is not a plain name. + pub facet: Option, +} + #[derive(Clone, Debug)] pub enum AtomicNarrowOp { Is(Expr), @@ -78,10 +122,22 @@ pub enum AtomicNarrowOp { NotHasAttr(Name), GetAttr(Name, Option>), NotGetAttr(Name, Option>), - TypeGuard(Type, Arguments), - NotTypeGuard(Type, Arguments), - TypeIs(Type, Arguments), - NotTypeIs(Type, Arguments), + /// A `TypeGuard`/`TypeIs` call resolved by standard call resolution. Standard + /// resolution identifies which actual argument supplies the guarded parameter + /// (`GuardNarrowing::subject`); only the candidate matching that argument contributes + /// a narrow, and the other candidates of the same call contribute nothing in either + /// polarity. + TypeGuard(Box), + NotTypeGuard(Box), + TypeIs(Box), + NotTypeIs(Box), + /// A call that may narrow the type of one of its arguments via a `TypeGuard`/`TypeIs` + /// result. Each candidate describes one actual argument: its keyword name (`None` for + /// the first positional argument) and the facet subject of its value expression, if + /// any. The call is shared because one call can produce candidates for several + /// subjects, and cloning the call AST for each of them is quadratic. + Call(Arc, Vec), + NotCall(Arc, Vec), // type(x) == y or type(x) is y TypeEq(Expr), TypeNotEq(Expr), @@ -116,9 +172,7 @@ pub enum AtomicNarrowOp { /// Used to narrow to mapping types (types that extend typing.Mapping) IsMapping, IsNotMapping, - /// (func, args) for a function call that may narrow the type of its first argument. - Call(Box, Arguments), - NotCall(Box, Arguments), + /// A narrow op applies to a name; these operations mean we are narrowing to the case /// when that name evaluates to a truthy or falsy value. IsTruthy, @@ -162,8 +216,8 @@ impl NarrowOp { | AtomicNarrowOp::IsNotSubclass(_) | AtomicNarrowOp::TypeEq(_) | AtomicNarrowOp::TypeNotEq(_) - | AtomicNarrowOp::Call(_, _) - | AtomicNarrowOp::NotCall(_, _) + | AtomicNarrowOp::Call(..) + | AtomicNarrowOp::NotCall(..) | AtomicNarrowOp::IsTruthy | AtomicNarrowOp::IsFalsy | AtomicNarrowOp::MatchSubjectElement(_, _) @@ -223,17 +277,37 @@ impl DisplayWith for AtomicNarrowOp { AtomicNarrowOp::NotGetAttr(attr, None) => { write!(f, "NotGetAttr({attr}, None)") } - AtomicNarrowOp::TypeGuard(t, arguments) => { - write!(f, "TypeGuard({t}, {})", arguments.display_with(ctx)) + AtomicNarrowOp::TypeGuard(narrowing) => { + write!( + f, + "TypeGuard({}, {})", + narrowing.func_ty, + narrowing.args.display_with(ctx) + ) } - AtomicNarrowOp::NotTypeGuard(t, arguments) => { - write!(f, "NotTypeGuard({t}, {})", arguments.display_with(ctx)) + AtomicNarrowOp::NotTypeGuard(narrowing) => { + write!( + f, + "NotTypeGuard({}, {})", + narrowing.func_ty, + narrowing.args.display_with(ctx) + ) } - AtomicNarrowOp::TypeIs(t, arguments) => { - write!(f, "TypeIs({t}, {})", arguments.display_with(ctx)) + AtomicNarrowOp::TypeIs(narrowing) => { + write!( + f, + "TypeIs({}, {})", + narrowing.func_ty, + narrowing.args.display_with(ctx) + ) } - AtomicNarrowOp::NotTypeIs(t, arguments) => { - write!(f, "NotTypeIs({t}, {})", arguments.display_with(ctx)) + AtomicNarrowOp::NotTypeIs(narrowing) => { + write!( + f, + "NotTypeIs({}, {})", + narrowing.func_ty, + narrowing.args.display_with(ctx) + ) } AtomicNarrowOp::TypeEq(expr) => write!(f, "TypeEq({})", expr.display_with(ctx)), AtomicNarrowOp::TypeNotEq(expr) => write!(f, "TypeNotEq({})", expr.display_with(ctx)), @@ -256,17 +330,17 @@ impl DisplayWith for AtomicNarrowOp { AtomicNarrowOp::IsNotSequence => write!(f, "IsNotSequence"), AtomicNarrowOp::IsMapping => write!(f, "IsMapping"), AtomicNarrowOp::IsNotMapping => write!(f, "IsNotMapping"), - AtomicNarrowOp::Call(expr, arguments) => write!( + AtomicNarrowOp::Call(call, _) => write!( f, "Call({}, {})", - expr.display_with(ctx), - arguments.display_with(ctx) + call.func.display_with(ctx), + call.args.display_with(ctx) ), - AtomicNarrowOp::NotCall(expr, arguments) => write!( + AtomicNarrowOp::NotCall(call, _) => write!( f, "NotCall({}, {})", - expr.display_with(ctx), - arguments.display_with(ctx) + call.func.display_with(ctx), + call.args.display_with(ctx) ), AtomicNarrowOp::IsTruthy => write!(f, "IsTruthy"), AtomicNarrowOp::IsFalsy => write!(f, "IsFalsy"), @@ -370,34 +444,34 @@ impl AtomicNarrowOp { Self::IsNotSequence => Some(format!("not isinstance({subject}, Sequence)")), Self::IsMapping => Some(format!("isinstance({subject}, Mapping)")), Self::IsNotMapping => Some(format!("not isinstance({subject}, Mapping)")), - Self::Call(expr, arguments) => { - let func = snippet(expr.range())?; - let args = snippet(arguments.range()).unwrap_or_default(); + Self::Call(call, _) => { + let func = snippet(call.func.range())?; + let args = snippet(call.args.range()).unwrap_or_default(); Some(format!("{func}{args}")) } - Self::NotCall(expr, arguments) => { - let func = snippet(expr.range())?; - let args = snippet(arguments.range()).unwrap_or_default(); + Self::NotCall(call, _) => { + let func = snippet(call.func.range())?; + let args = snippet(call.args.range()).unwrap_or_default(); Some(format!("not {func}{args}")) } Self::IsTruthy => Some(subject.to_owned()), Self::IsFalsy => Some(format!("not {subject}")), Self::PolarsColumnMutation(_) => None, - Self::TypeGuard(_, arguments) => Some(format!( + Self::TypeGuard(narrowing) => Some(format!( "TypeGuard{}", - snippet(arguments.range()).unwrap_or_default() + snippet(narrowing.args.range()).unwrap_or_default() )), - Self::NotTypeGuard(_, arguments) => Some(format!( + Self::NotTypeGuard(narrowing) => Some(format!( "not TypeGuard{}", - snippet(arguments.range()).unwrap_or_default() + snippet(narrowing.args.range()).unwrap_or_default() )), - Self::TypeIs(_, arguments) => Some(format!( + Self::TypeIs(narrowing) => Some(format!( "TypeIs{}", - snippet(arguments.range()).unwrap_or_default() + snippet(narrowing.args.range()).unwrap_or_default() )), - Self::NotTypeIs(_, arguments) => Some(format!( + Self::NotTypeIs(narrowing) => Some(format!( "not TypeIs{}", - snippet(arguments.range()).unwrap_or_default() + snippet(narrowing.args.range()).unwrap_or_default() )), Self::Placeholder => None, Self::MatchSubjectElement(_, _) => None, @@ -405,6 +479,29 @@ impl AtomicNarrowOp { } } + /// The type the subject narrows to. + pub fn narrowed_type(&self) -> &Type { + match self { + Self::TypeIs(narrowing) + | Self::NotTypeIs(narrowing) + | Self::TypeGuard(narrowing) + | Self::NotTypeGuard(narrowing) => &narrowing.narrowed, + _ => unreachable!("guard narrowing accessor on non-guard op"), + } + } + + /// The guarded argument identified by standard call resolution for this narrowing + /// call, when it identified one. + pub fn guard_subject(&self) -> Option<&GuardSubject> { + match self { + Self::TypeIs(narrowing) + | Self::NotTypeIs(narrowing) + | Self::TypeGuard(narrowing) + | Self::NotTypeGuard(narrowing) => narrowing.subject.as_ref(), + _ => None, + } + } + pub fn negate(&self) -> Self { match self { Self::Is(v) => Self::IsNot(v.clone()), @@ -435,14 +532,14 @@ impl AtomicNarrowOp { Self::IsNotSequence => Self::IsSequence, Self::IsMapping => Self::IsNotMapping, Self::IsNotMapping => Self::IsMapping, - Self::TypeGuard(ty, args) => Self::NotTypeGuard(ty.clone(), args.clone()), - Self::NotTypeGuard(ty, args) => Self::TypeGuard(ty.clone(), args.clone()), - Self::TypeIs(ty, args) => Self::NotTypeIs(ty.clone(), args.clone()), - Self::NotTypeIs(ty, args) => Self::TypeIs(ty.clone(), args.clone()), + Self::TypeGuard(narrowing) => Self::NotTypeGuard(narrowing.clone()), + Self::NotTypeGuard(narrowing) => Self::TypeGuard(narrowing.clone()), + Self::TypeIs(narrowing) => Self::NotTypeIs(narrowing.clone()), + Self::NotTypeIs(narrowing) => Self::TypeIs(narrowing.clone()), Self::TypeEq(v) => Self::TypeNotEq(v.clone()), Self::TypeNotEq(v) => Self::TypeEq(v.clone()), - Self::Call(f, args) => Self::NotCall(f.clone(), args.clone()), - Self::NotCall(f, args) => Self::Call(f.clone(), args.clone()), + Self::Call(call, candidates) => Self::NotCall(call.clone(), candidates.clone()), + Self::NotCall(call, candidates) => Self::Call(call.clone(), candidates.clone()), Self::IsTruthy => Self::IsFalsy, Self::IsFalsy => Self::IsTruthy, Self::PolarsColumnMutation(kind) => Self::PolarsColumnMutation(kind.clone()), @@ -915,6 +1012,53 @@ impl NarrowOps { } } + /// Records one candidate argument for a narrowing call. + /// + /// Candidates whose value expressions share a subject root share the entry for that + /// root: the answers phase selects the candidate that standard call resolution + /// identifies as the guarded argument, and the other candidates of the same call + /// contribute nothing in either polarity. Candidates of different calls on the same + /// subject are separate predicates and combine with AND. + fn add_call_candidate( + &mut self, + call: &Arc, + kw: Option, + value: &Expr, + range: TextRange, + ) { + for subject in expr_to_subjects(value) { + let (name, facet) = match subject { + NarrowingSubject::Name(name) => (name, None), + NarrowingSubject::Facets(name, facets) => (name, Some(facets)), + }; + let candidate = CallCandidate { + kw: kw.clone(), + facet, + }; + match self.0.entry(name) { + Entry::Occupied(mut entry) => { + let (op, _) = entry.get_mut(); + if let NarrowOp::Atomic(_, AtomicNarrowOp::Call(existing_call, candidates)) = op + && Arc::ptr_eq(existing_call, call) + { + candidates.push(candidate); + } else { + op.and(NarrowOp::Atomic( + None, + AtomicNarrowOp::Call(call.clone(), vec![candidate]), + )); + } + } + Entry::Vacant(entry) => { + entry.insert(( + NarrowOp::Atomic(None, AtomicNarrowOp::Call(call.clone(), vec![candidate])), + range, + )); + } + } + } + } + pub fn negate(&self) -> Self { Self( self.0 @@ -1331,14 +1475,31 @@ impl NarrowOps { func, arguments: args @ Arguments { args: posargs, .. }, }, - ) if !posargs.is_empty() => { - // This may be a function call that narrows the type of its first argument. Record - // it as a possible narrowing operation that we'll resolve in the answers phase. - Self::from_single_narrow_op( - &posargs[0], - AtomicNarrowOp::Call(Box::new((**func).clone()), args.clone()), - test.range(), - ) + ) if !posargs.is_empty() || !args.keywords.is_empty() => { + // This may be a function call whose `TypeGuard`/`TypeIs` result narrows one of + // its arguments. Standard call resolution identifies the guarded argument in + // the answers phase; record one candidate per actual argument here. Candidates + // whose value expressions share a subject root share the entry for that root, + // so a call contributes at most one narrow per subject in either polarity. + let call = Arc::new(NarrowingCall { + func: Box::new((**func).clone()), + args: args.clone(), + }); + let mut narrow_ops = Self::new(); + if let Some(first) = posargs.first() { + narrow_ops.add_call_candidate(&call, None, first, test.range()); + } + for keyword in &args.keywords { + if let Some(kw_name) = &keyword.arg { + narrow_ops.add_call_candidate( + &call, + Some(kw_name.id.clone()), + &keyword.value, + test.range(), + ); + } + } + narrow_ops } Expr::Named(named) => { let mut target_narrow = Self::from_single_narrow_op( diff --git a/pyrefly/lib/test/narrow.rs b/pyrefly/lib/test/narrow.rs index b214ddd662..6263e81dc3 100644 --- a/pyrefly/lib/test/narrow.rs +++ b/pyrefly/lib/test/narrow.rs @@ -1676,6 +1676,194 @@ def f(x: Cat | Dog): "#, ); +testcase!( + test_p3_minimal, + r#" +from typing import assert_type, TypeIs + + +def is_str_on_value(value: object, other_parameter: object) -> TypeIs[str]: + return isinstance(value, str) + + +def both_args_are_value(value: str | int) -> None: + if is_str_on_value(value=value, other_parameter=value): + assert_type(value, str) + else: + assert_type(value, int) + "#, +); + +testcase!( + test_typeguard_keyword_argument, + r#" +from typing import TypeGuard, TypeIs, assert_type, overload + +def is_str(x: str | None) -> TypeGuard[str]: + return x is not None + +def is_str_is(x: str | None) -> TypeIs[str]: + return x is not None + +def need_str(x: str) -> None: ... + +def kw_guard(x: str | None) -> None: + if is_str(x=x): + assert_type(x, str) + else: + assert_type(x, str | None) + +def kw_is(x: str | None) -> None: + if is_str_is(x=x): + need_str(x) + else: + assert_type(x, None) + +def kw_is_not(x: str | None) -> None: + if not is_str_is(x=x): + assert_type(x, None) + else: + need_str(x) + +def kw_not(x: str | None) -> None: + if not is_str(x=x): + need_str(x) # E: Argument `str | None` is not assignable to parameter `x` with type `str` in function `need_str` + else: + assert_type(x, str) + +# A keyword that binds a later parameter does not narrow that keyword's value. +def guard2(y: object, x: str | None) -> TypeGuard[str]: + return x is not None + +def kw_second(y: object, x: str | None) -> None: + if guard2(y, x=x): + need_str(y) + need_str(x) # E: Argument `str | None` is not assignable to parameter `x` with type `str` in function `need_str` + +# A keyword cannot bind a positional-only first parameter. +def guard_po(x: str | None, /) -> TypeGuard[str]: + return x is not None + +def kw_po(x: str | None) -> None: + if guard_po(x=x): # E: Expected argument `x` to be positional + need_str(x) # E: Argument `str | None` is not assignable to parameter `x` with type `str` in function `need_str` + +# Methods: the bound receiver is skipped when finding the guarded parameter. +class Guards: + def is_str(self, x: str | None) -> TypeGuard[str]: + return x is not None + +def kw_method(g: Guards, x: str | None) -> None: + if g.is_str(x=x): + need_str(x) + +# An unbound method call supplies the receiver explicitly. The receiver is never the +# argument the guard result constrains, so it is not narrowed with either spelling, and +# the value argument is narrowed through its keyword. +def kw_unbound(g: Guards, x: str | None) -> None: + if Guards.is_str(self=g, x=x): + g.upper() # E: Object of class `Guards` has no attribute `upper` + need_str(x) + assert_type(g, Guards) + +def kw_unbound_positional(g: Guards, x: str | None) -> None: + if Guards.is_str(g, x): + g.upper() # E: Object of class `Guards` has no attribute `upper` + assert_type(g, Guards) + +# Overloads: the guarded argument comes from the signature the call resolves to. +@overload +def oguard(x: int) -> TypeGuard[int]: ... +@overload +def oguard(x: str | None) -> TypeGuard[str]: ... +def oguard(x: int | str | None) -> bool: + return x is not None + +def kw_overload(x: str | None) -> None: + if oguard(x=x): + need_str(x) + +# Attributes narrow like facets. +class Box: + field: str | None + +def kw_facet(b: Box) -> None: + if is_str(x=b.field): + need_str(b.field) + +# Loops and asserts work like `if`. +def kw_while(x: str | None) -> None: + while is_str(x=x): + need_str(x) + +def kw_assert(x: str | None) -> None: + assert is_str(x=x) + need_str(x) + "#, +); + +// The guarded argument comes from the signature the call resolves to; the other +// signature's subject is untouched. +testcase!( + test_typeguard_keyword_overload_subject, + r#" +from typing import Literal, TypeGuard, assert_type, overload + +@overload +def guard(a: object, b: object, *, mode: Literal["a"]) -> TypeGuard[str]: ... +@overload +def guard(b: object, a: object, *, mode: Literal["b"]) -> TypeGuard[int]: ... +def guard(*args: object, **kwargs: object) -> bool: + key = "a" if kwargs.get("mode") == "a" else "b" + value = kwargs[key] if key in kwargs else args[0] + return isinstance(value, str if key == "a" else int) + +def unsafe(a: object, b: object) -> None: + if guard(a=a, b=b, mode="b"): + assert_type(a, object) + assert_type(b, int) + "#, +); + +testcase!( + test_typeguard_keyword_single_narrow_per_call, + r#" +from typing import TypeGuard, TypeIs, assert_type + +def is_str_on_value(value: object, other_parameter: object) -> TypeIs[str]: + return isinstance(value, str) + +def both_args_are_value(value: str | int) -> None: + # `value` supplies both the guarded parameter and another parameter. The call is a + # single predicate, so both branches narrow (the old AND-of-candidates code lost the + # negative narrowing here). + if is_str_on_value(value=value, other_parameter=value): + assert_type(value, str) + else: + assert_type(value, int) + + "#, +); + +testcase!( + test_typeguard_keyword_overload_disagree_subject, + r#" +from typing import TypeGuard, assert_type, overload + +@overload +def guard(a: object, b: object) -> TypeGuard[str]: ... +@overload +def guard(b: object, a: object) -> TypeGuard[str]: ... +def guard(*args: object, **kwargs: object) -> bool: + raise AssertionError + +def ambiguous(a: object, b: object) -> None: + if guard(a=a, b=b): + assert_type(a, str) + assert_type(b, object) + "#, +); + testcase!( test_typeis, r#"