diff --git a/crates/pyrefly_types/src/function.rs b/crates/pyrefly_types/src/function.rs index b7e88ebc5a..f56dab6070 100644 --- a/crates/pyrefly_types/src/function.rs +++ b/crates/pyrefly_types/src/function.rs @@ -196,10 +196,12 @@ pub struct FuncFlags { /// Parameter indices whose annotations directly name a type parameter of the defining class. /// This is used when a subclass binds that parameter to a shape `Flag`. pub shape_flag_constructor_sources: Option>>, - /// A function decorated with `@deprecated` - pub deprecation: Option, + /// A function decorated with `@deprecated`. Boxing keeps callable types small. + pub deprecation: Option>, /// Metadata for `@property`, `@foo.setter`, and `@foo.deleter`. pub property_metadata: Option, + /// The object referenced by `__wrapped__`, without changing the wrapper's signature. + pub wrapped: Option>, /// A function decorated with `functools.cached_property` or equivalent. pub is_cached_property: bool, pub has_enum_member_decoration: bool, @@ -436,6 +438,9 @@ pub enum FunctionKind { AssertType, AssertShape, RevealType, + InspectUnwrap, + FunctoolsWraps, + FunctoolsUpdateWrapper, Final, RuntimeCheckable, Synthesized(Arc), @@ -525,6 +530,9 @@ impl FunctionKind { ("typing" | "typing_extensions", "assert_type") => Self::AssertType, ("shape_extensions", "assert_shape") => Self::AssertShape, ("typing" | "typing_extensions", "reveal_type") => Self::RevealType, + ("inspect", "unwrap") => Self::InspectUnwrap, + ("functools", "wraps") => Self::FunctoolsWraps, + ("functools", "update_wrapper") => Self::FunctoolsUpdateWrapper, ("typing" | "typing_extensions", "final") => Self::Final, ("typing" | "typing_extensions", "runtime_checkable") => Self::RuntimeCheckable, ("typing" | "typing_extensions", "dataclass_transform") => Self::DataclassTransform, @@ -564,6 +572,9 @@ impl FunctionKind { Self::AssertType => ModuleName::typing(), Self::AssertShape => ModuleName::from_str("shape_extensions"), Self::RevealType => ModuleName::typing(), + Self::InspectUnwrap => ModuleName::from_str("inspect"), + Self::FunctoolsWraps => ModuleName::functools(), + Self::FunctoolsUpdateWrapper => ModuleName::functools(), Self::RuntimeCheckable => ModuleName::typing(), Self::CallbackProtocol(cls) => cls.qname().module_name(), Self::SingleDispatchRegister(_) => ModuleName::functools(), @@ -605,6 +616,9 @@ impl FunctionKind { Self::AssertType => Cow::Owned(Name::new_static("assert_type")), Self::AssertShape => Cow::Owned(Name::new_static("assert_shape")), Self::RevealType => Cow::Owned(Name::new_static("reveal_type")), + Self::InspectUnwrap => Cow::Owned(Name::new_static("unwrap")), + Self::FunctoolsWraps => Cow::Owned(Name::new_static("wraps")), + Self::FunctoolsUpdateWrapper => Cow::Owned(Name::new_static("update_wrapper")), Self::RuntimeCheckable => Cow::Owned(Name::new_static("runtime_checkable")), Self::CallbackProtocol(_) => Cow::Owned(dunder::CALL), Self::SingleDispatchRegister(_) => Cow::Owned(Name::new_static("register")), @@ -646,6 +660,9 @@ impl FunctionKind { Self::AssertType => None, Self::AssertShape => None, Self::RevealType => None, + Self::InspectUnwrap => None, + Self::FunctoolsWraps => None, + Self::FunctoolsUpdateWrapper => None, Self::RuntimeCheckable => None, Self::NumbaJit => None, Self::NumbaNjit => None, diff --git a/crates/pyrefly_types/src/types.rs b/crates/pyrefly_types/src/types.rs index 7f6d238612..9f4917812c 100644 --- a/crates/pyrefly_types/src/types.rs +++ b/crates/pyrefly_types/src/types.rs @@ -1415,7 +1415,17 @@ impl Type { *found = true; return; } - ty.recurse(&mut |child| walk(child, found)); + // Unknown types in `__wrapped__` do not make the wrapper's signature unknown. + if ty + .toplevel_func_metadata() + .is_some_and(|metadata| metadata.flags.wrapped.is_some()) + { + for (signature, _) in ty.toplevel_callable_signatures() { + signature.visit(&mut |child| walk(child, found)); + } + } else { + ty.recurse(&mut |child| walk(child, found)); + } } let mut found = false; walk(self, &mut found); @@ -2033,7 +2043,7 @@ impl Type { pub fn function_deprecation(&self) -> Option<&Deprecation> { self.toplevel_func_metadata() - .and_then(|meta| meta.flags.deprecation.as_ref()) + .and_then(|meta| meta.flags.deprecation.as_deref()) } pub fn has_final_decoration(&self) -> bool { diff --git a/pyrefly/lib/alt/attr.rs b/pyrefly/lib/alt/attr.rs index c7b41e3a42..4bce1c2f73 100644 --- a/pyrefly/lib/alt/attr.rs +++ b/pyrefly/lib/alt/attr.rs @@ -494,6 +494,8 @@ enum AttributeBase1 { EnumLiteral(LitEnum), LiteralString, ClassInstance(ClassType), + /// A function with a precise `__wrapped__` attribute and ordinary function behavior. + WrappedFunction(ClassType, Type), ClassObject(ClassBase), Module(ModuleType), /// Attribute access on a type parameter in a value position. @@ -744,6 +746,10 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { AttributeBase1::ClassInstance(class) => { self.add_class_fields(class.class_object(), candidates); } + AttributeBase1::WrappedFunction(class, _) => { + self.add_class_fields(class.class_object(), candidates); + candidates.insert(Name::new_static("__wrapped__")); + } AttributeBase1::ClassObject(class_base) => { self.add_class_fields(class_base.class_object(), candidates); } @@ -1941,6 +1947,17 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { self.lookup_attr_static1((**protocol_base).clone(), attr_name, acc) } } + AttributeBase1::WrappedFunction(class, wrapped) => { + if attr_name.as_str() == "__wrapped__" { + acc.found_type(wrapped.clone(), base); + } else { + self.lookup_attr_static1( + AttributeBase1::ClassInstance(class.clone()), + attr_name, + acc, + ); + } + } AttributeBase1::BoundMethod(bound_func) => { let method_type_base = AttributeBase1::ClassInstance(self.stdlib.method_type().clone()); @@ -2293,6 +2310,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { } } AttributeBase1::ClassInstance(cls) + | AttributeBase1::WrappedFunction(cls, _) | AttributeBase1::SelfType(cls) | AttributeBase1::EnumLiteral(LitEnum { class: cls, .. }) | AttributeBase1::Quantified(_, cls) @@ -2738,6 +2756,17 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { Type::Callable(c) if matches!(c.params, Params::Partial(_)) => acc.push( AttributeBase1::ClassInstance(self.stdlib.partial(c.ret.clone())), ), + _ if !matches!(ty, Type::BoundMethod(_)) + && let Some(metadata) = ty.toplevel_func_metadata() + && let Some(wrapped) = &metadata.flags.wrapped => + { + let class = if let FunctionKind::CallbackProtocol(cls) = &metadata.kind { + (**cls).clone() + } else { + self.stdlib.function_type().clone() + }; + acc.push(AttributeBase1::WrappedFunction(class, (**wrapped).clone())); + } Type::Callable(_) => acc.push(AttributeBase1::ClassInstance( self.stdlib.function_type().clone(), )), @@ -3587,6 +3616,20 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { | AttributeBase1::Quantified(_, class) => { self.completions_class_type(class, expected_attribute_name, res) } + AttributeBase1::WrappedFunction(class, _) => { + self.completions_class_type(class, expected_attribute_name, res); + let name = Name::new_static("__wrapped__"); + if expected_attribute_name.is_none_or(|expected| *expected == name) { + res.push(AttrInfo { + name, + ty: None, + is_deprecated: false, + definition: AttrDefinition::Synthetic, + is_reexport: false, + is_from_object: false, + }); + } + } AttributeBase1::ShapedArrayInstance(tensor) => { self.completions_class_type(&tensor.base_class, expected_attribute_name, res) } diff --git a/pyrefly/lib/alt/call.rs b/pyrefly/lib/alt/call.rs index d9624854f4..c3f30ee7d1 100644 --- a/pyrefly/lib/alt/call.rs +++ b/pyrefly/lib/alt/call.rs @@ -1965,6 +1965,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { None } }; + let wrapped_arg = self.functools_wrapped_arg(metadata, args, keywords, errors); let outcome = match call_target { CallTarget::Class(cls, constructor_kind, as_quantified_bound) => { if cls.has_qname("typing", "Any") { @@ -2226,7 +2227,18 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { ty: res, overload_table, } = outcome; - let res = if let Some(func_metadata) = kw_metadata { + let res = if let Some((func_metadata, wrapped)) = wrapped_arg.as_ref() + && func_metadata.kind == FunctionKind::FunctoolsWraps + { + let mut kws = TypeMap::new(); + kws.0 + .insert(Name::new_static("__wrapped__"), wrapped.clone()); + self.heap.mk_kw_call(KwCall { + func_metadata: func_metadata.clone(), + keywords: kws, + return_ty: res, + }) + } else if let Some(func_metadata) = kw_metadata { // The call form `dataclass(C)` transforms `C` in place, so reject the same // class kinds as the `@dataclass` decorator (see `report_forbidden_dataclass_target`). // The decorator path never reaches here: a bare decorator is not a call. @@ -2266,6 +2278,19 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { } else { res }; + // Keep ordinary functions as functions so descriptor binding and function flags survive. + let res = if let Some((metadata, wrapped)) = wrapped_arg + && metadata.kind == FunctionKind::FunctoolsUpdateWrapper + && let Some(mut wrapper) = self + .positional_arg_type(args, 0, errors) + .or_else(|| self.keyword_arg_type(keywords, "wrapper", errors)) + && let Some(metadata) = wrapper.toplevel_func_metadata_mut() + { + metadata.flags.wrapped = Some(Box::new(wrapped)); + wrapper + } else { + res + }; CallOutcome { ty: res, overload_table, @@ -2524,6 +2549,49 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { CallOutcome::of_ty(self.intersect_guarantees(&results)) } + fn functools_wrapped_arg( + &self, + metadata: Option<&FuncMetadata>, + args: &[CallArg], + keywords: &[CallKeyword], + errors: &ErrorCollector, + ) -> Option<(FuncMetadata, Type)> { + let metadata = metadata?; + let wrapped = match metadata.kind { + FunctionKind::FunctoolsWraps => self + .positional_arg_type(args, 0, errors) + .or_else(|| self.keyword_arg_type(keywords, "wrapped", errors)), + FunctionKind::FunctoolsUpdateWrapper => self + .positional_arg_type(args, 1, errors) + .or_else(|| self.keyword_arg_type(keywords, "wrapped", errors)), + _ => None, + }?; + Some((metadata.clone(), wrapped)) + } + + fn positional_arg_type( + &self, + args: &[CallArg], + index: usize, + errors: &ErrorCollector, + ) -> Option { + match args.get(index)? { + CallArg::Arg(value) => Some((*value).infer(self, errors)), + CallArg::Star(..) => None, + } + } + + fn keyword_arg_type( + &self, + keywords: &[CallKeyword], + name: &str, + errors: &ErrorCollector, + ) -> Option { + keywords.iter().find_map(|kw| { + (kw.arg.is_some_and(|id| id.as_str() == name)).then(|| kw.value.infer(self, errors)) + }) + } + /// Wrapper for `callable_infer` that handles trying a call with and without a contextual hint. fn call_infer_inner( &self, @@ -3083,6 +3151,11 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { errors, ) } + Some(CalleeKind::Function(FunctionKind::InspectUnwrap)) + if let Some(ret) = self.call_inspect_unwrap(x, errors) => + { + ret + } // `attr.evolve` validates kwargs like `dataclasses.replace`; `attr.assoc` validates // against attribute names, including `init=False` fields. Both require an attrs class. Some(CalleeKind::Function( @@ -3368,6 +3441,65 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { } } + fn call_inspect_unwrap(&self, x: &ExprCall, errors: &ErrorCollector) -> Option { + if x.arguments.args.len() != 1 || matches!(x.arguments.args[0], Expr::Starred(_)) { + return None; + } + if !x.arguments.keywords.iter().all(|kw| { + kw.arg.as_ref().is_some_and(|id| id.as_str() == "stop") + && matches!(kw.value, Expr::NoneLiteral(_)) + }) { + return None; + } + + let arg = &x.arguments.args[0]; + let ty = self.expr_infer(arg, errors); + Some(self.inspect_unwrapped_type(&ty, arg.range(), &mut SmallSet::new(), false)) + } + + /// Follow each wrapped chain independently, retaining uncertainty about hidden attributes. + fn inspect_unwrapped_type( + &self, + ty: &Type, + range: TextRange, + seen: &mut SmallSet, + followed_wrapped: bool, + ) -> Type { + self.distribute_over_union(ty, |ty| { + if !seen.insert(ty.clone()) { + return self.heap.mk_any_explicit(); + } + let wrapped_attr = Name::new_static("__wrapped__"); + let result = if !self.has_attr(ty, &wrapped_attr) { + // A callable annotation or an in-place mutation can hide `__wrapped__`. + // Preserve only a terminal function reached through a known wrapped edge. + if followed_wrapped && ty.toplevel_func_metadata().is_some() { + ty.clone() + } else { + self.heap.mk_any_explicit() + } + } else { + let attr_errors = self.error_swallower(); + let next = self.type_of_attr_get( + ty, + &wrapped_attr, + range, + &attr_errors, + ErrorKind::MissingAttribute, + None, + "__wrapped__ lookup for inspect.unwrap", + ); + if attr_errors.is_empty() { + self.inspect_unwrapped_type(&next, range, seen, true) + } else { + self.heap.mk_any_explicit() + } + }; + seen.shift_remove(ty); + result + }) + } + pub fn freeform_call_infer( &self, ty: Type, diff --git a/pyrefly/lib/alt/function.rs b/pyrefly/lib/alt/function.rs index 3d91c68246..c72f5e92de 100644 --- a/pyrefly/lib/alt/function.rs +++ b/pyrefly/lib/alt/function.rs @@ -1160,7 +1160,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { true } SpecialDecorator::Deprecated(deprecation) => { - flags.deprecation = Some((**deprecation).clone()); + flags.deprecation = Some(Box::new((**deprecation).clone())); true } SpecialDecorator::PropertySetter(decorator) => { @@ -1996,6 +1996,16 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { { return decoratee; } + let wrapped = if let Type::KwCall(call) = &decorator + && call.has_function_kind(FunctionKind::FunctoolsWraps) + { + call.keywords + .0 + .get(&Name::new_static("__wrapped__")) + .cloned() + } else { + None + }; let application = self.prepare_decorator_application(decorator, decoratee, range, errors); // Run a decorator call, buffering errors so we can decide between the primary // and Self-rewritten fallback without double-reporting. @@ -2032,8 +2042,13 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { errors.extend(primary_errors); primary_return }; - let decorated_value = + let mut decorated_value = self.decorate_returned_callable(raw_return, metadata, &application.original_decoratee); + if let Some(wrapped) = wrapped + && let Some(metadata) = decorated_value.toplevel_func_metadata_mut() + { + metadata.flags.wrapped = Some(Box::new(wrapped)); + } // Given the raw `decorated_value`, which may include `Type::Quantified` type variables // coming from a `Forall` in the original decoratee, we need to create the proper output // type: diff --git a/pyrefly/lib/alt/special_calls.rs b/pyrefly/lib/alt/special_calls.rs index 054498dc51..bc9a550505 100644 --- a/pyrefly/lib/alt/special_calls.rs +++ b/pyrefly/lib/alt/special_calls.rs @@ -865,6 +865,20 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> { kws: &[CallKeyword], errors: &ErrorCollector, ) -> Option { + if let Type::KwCall(call) = callee + && call.has_function_kind(FunctionKind::FunctoolsWraps) + && kws.is_empty() + && let [CallArg::Arg(arg)] = args + && let Some(wrapped) = call.keywords.0.get(&Name::new_static("__wrapped__")) + { + let arg_errors = self.error_collector(); + let mut wrapper = arg.infer(self, &arg_errors); + if let Some(metadata) = wrapper.toplevel_func_metadata_mut() { + metadata.flags.wrapped = Some(Box::new(wrapped.clone())); + errors.extend(arg_errors); + return Some(wrapper); + } + } let decorator = Decorator { ty: callee.clone(), deprecation: None, diff --git a/pyrefly/lib/solver/subset.rs b/pyrefly/lib/solver/subset.rs index e203027570..53c3f24062 100644 --- a/pyrefly/lib/solver/subset.rs +++ b/pyrefly/lib/solver/subset.rs @@ -1955,6 +1955,8 @@ impl<'solver, 'subset, Ans: LookupAnswer> Subset<'solver, 'subset, Ans> { (_, Type::ClassType(want)) if want.is_builtin("object") => { Ok(()) // everything is an instance of `object` } + (Type::KwCall(got), _) => self.is_subset_eq(&got.return_ty, want), + (_, Type::KwCall(want)) => self.is_subset_eq(got, &want.return_ty), (got_ty, want_ty) if let Some(got_alias) = as_type_alias(got_ty) && let Some(want_alias) = as_type_alias(want_ty) => diff --git a/pyrefly/lib/test/decorators.rs b/pyrefly/lib/test/decorators.rs index 99904106c8..8a5fb5d8f6 100644 --- a/pyrefly/lib/test/decorators.rs +++ b/pyrefly/lib/test/decorators.rs @@ -53,6 +53,239 @@ reveal_type(decorated) # E: revealed type: (x: int) -> str "#, ); +testcase!( + test_inspect_unwrap_follows_wrapped_chain, + r#" +from functools import wraps +from inspect import unwrap +from typing import assert_type, reveal_type + +def inner(x: int) -> str: + return f"{x}" + +@wraps(inner) +def middle(x: object) -> None: + pass + +@wraps(middle) +def outer(x: object) -> None: + pass + +reveal_type(unwrap(outer)) # E: revealed type: (x: int) -> str +assert_type(outer(object()), None) +assert_type(middle.__wrapped__(1), str) +assert_type(outer.__wrapped__(object()), None) +assert_type(unwrap(outer, stop=None)(1), str) + "#, +); + +testcase!( + test_functools_wrapped_call_chains, + r#" +from functools import update_wrapper, wraps +from inspect import unwrap +from typing import assert_type + +def inner(x: int) -> str: + return str(x) + +def middle(x: object) -> None: + pass + +def outer(x: object) -> bool: + return False + +wrapped_middle = update_wrapper(middle, inner) +wrapped_outer = update_wrapper(wrapper=outer, wrapped=wrapped_middle) +assert_type(wrapped_outer(object()), bool) +assert_type(wrapped_outer.__wrapped__(object()), None) +assert_type(unwrap(wrapped_outer)(1), str) + +called_middle = wraps(inner)(middle) +called_outer = wraps(called_middle)(outer) +assert_type(called_outer(object()), bool) +assert_type(called_outer.__wrapped__(object()), None) +assert_type(unwrap(called_outer)(1), str) + "#, +); + +testcase!( + test_functools_wraps_generic_chain, + r#" +from functools import wraps +from inspect import unwrap +from typing import assert_type + +def inner[T](x: T) -> list[T]: + return [x] + +@wraps(inner) +def outer[T](x: T) -> tuple[T]: + return (x,) + +assert_type(outer(1), tuple[int]) +assert_type(outer.__wrapped__(1), list[int]) +assert_type(unwrap(outer)(1), list[int]) + "#, +); + +testcase!( + test_functools_wraps_preserves_method_and_property, + r#" +from functools import wraps +from typing import Any, Callable, assert_type + +def decorate(fn: Callable[[Any, int], int]): + @wraps(fn) + def wrapper(self: Any, x: int) -> int: + return fn(self, x) + return wrapper + +def property_source(self: Any) -> int: + return 1 + +class Parent: + def method(self, x: int) -> int: + return x + +class Child(Parent): + @decorate + def method(self, x: int) -> int: + return x + + @property + @wraps(property_source) + def value(self) -> int: + return 1 + +assert_type(Child().method(1), int) +assert_type(Child().value, int) + "#, +); + +testcase!( + test_functools_wraps_unknown_target_preserves_signature, + TestEnv::new().enable_partially_unknown_type_error(), + r#" +from functools import wraps +from typing import Callable, assert_type + +def untyped(x): + return x + +@wraps(untyped) # E: The type of this argument is partially unknown # E: The type of this argument is partially unknown +def wrapper(x: int) -> int: + return x + +def consume(f: Callable[[int], int]) -> int: + return f(1) + +assert_type(consume(wrapper), int) + "#, +); + +testcase!( + test_functools_update_wrapper_loop, + r#" +from functools import partial, update_wrapper +from typing import Callable + +def build(handler: Callable[[int], int], apps: tuple[list[Callable[..., int]], ...]) -> Callable[[int], int]: + for app in apps[::-1]: + for middleware in app: + handler = update_wrapper(partial(middleware, handler=handler), handler) + return handler + "#, +); + +testcase!( + test_inspect_unwrap_unknown_wrapping_state, + r#" +from functools import update_wrapper, wraps +from inspect import unwrap +from typing import Any, Callable, assert_type + +def inner() -> str: + return "wrapped" + +@wraps(inner) +def wrapper() -> int: + return 1 + +def annotated(fn: Callable[[], int]) -> None: + assert_type(unwrap(fn), Any) + + @wraps(fn) + def outer() -> bool: + return False + + assert_type(unwrap(outer), Any) + +annotated(wrapper) + +def mutated() -> int: + return 1 + +update_wrapper(mutated, inner) +assert_type(unwrap(mutated), Any) + "#, +); + +testcase!( + test_inspect_unwrap_union_branches, + r#" +from functools import wraps +from inspect import unwrap +from typing import Any, assert_type + +def inner() -> str: + return "wrapped" + +@wraps(inner) +def wrapper() -> int: + return 1 + +def plain() -> int: + return 2 + +@wraps(inner) +def another() -> bool: + return False + +def check(flag: bool) -> None: + choice = wrapper if flag else plain + assert_type(unwrap(choice)(), Any | str) + + @wraps(choice) + def outer() -> bool: + return False + + assert_type(unwrap(outer)(), int | str) + + shared_target = wrapper if flag else another + assert_type(unwrap(shared_target)(), str) + "#, +); + +testcase!( + test_inspect_unwrap_starred_arguments, + r#" +from functools import wraps +from inspect import unwrap +from typing import Any, assert_type + +def inner() -> str: + return "wrapped" + +@wraps(inner) +def wrapper() -> int: + return 1 + +assert_type(unwrap(*(wrapper,)), Any) +assert_type(unwrap(*(wrapper,), stop=None), Any) + "#, +); + testcase!( test_signature_modifying_function_decorator, r#"