fix support inspect unwrap to return precise type #3696

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<Box<Vec<usize>>>,
-    /// A function decorated with `@deprecated`
-    pub deprecation: Option<Deprecation>,
+    /// A function decorated with `@deprecated`. Boxing keeps callable types small.
+    pub deprecation: Option<Box<Deprecation>>,
     /// Metadata for `@property`, `@foo.setter`, and `@foo.deleter`.
     pub property_metadata: Option<PropertyMetadata>,
+    /// The object referenced by `__wrapped__`, without changing the wrapper's signature.
+    pub wrapped: Option<Box<Type>>,
     /// 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<FuncSymbol>),
@@ -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,

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 {

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)
             }

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<Type> {
+        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<Type> {
+        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<Type> {
+        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<Type>,
+        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,

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: