fix

Author: asukaminato0721

crates/pyrefly_config/src/error_kind.rs

@@ -295,6 +295,8 @@ pub enum ErrorKind {
     InvalidInheritance,
     /// Attempting to use a value that is not a valid kind of Literal.
     InvalidLiteral,
+    /// Expected arguments in a mock assertion do not match the patched target's signature.
+    InvalidMockAssertion,
     /// An error caused by incorrect usage of the @overload decorator.
     /// e.g. not defining multiple variants for an overloaded function.
     InvalidOverload,
@@ -664,6 +666,7 @@ impl ErrorKind {
             ErrorKind::InvalidAbstractMethod => Severity::Ignore,
             ErrorKind::InvalidCast => Severity::Ignore,
             ErrorKind::InvalidDecorator => Severity::Warn,
+            ErrorKind::InvalidMockAssertion => Severity::Ignore,
             ErrorKind::MethodAssign => Severity::Ignore,
             ErrorKind::MisplacedIgnore => Severity::Warn,
             ErrorKind::MissingAttributePatchTarget => Severity::Warn,

crates/pyrefly_types/src/display.rs

@@ -1107,6 +1107,9 @@ impl<'a> TypeDisplayContext<'a> {
                 // Display as the class name (e.g., MaxPool2d)
                 self.fmt_helper_generic(&Type::ClassType(module.class.clone()), false, output)
             }
+            Type::UnittestMock(mock) => {
+                self.fmt_helper_generic(&Type::ClassType(mock.class.clone()), false, output)
+            }
             Type::DataFrame(schema) if self.strip_library_schemas => {
                 self.fmt_helper_generic(&schema.underlying_type(), false, output)
             }

crates/pyrefly_types/src/function.rs

@@ -244,6 +244,8 @@ pub struct FuncFlags {
     /// A method directly inside a `Protocol` class.
     pub is_in_protocol_class: bool,
     pub is_in_type_checking_block: bool,
+    /// A mock assertion checks expected arguments without invoking the target function.
+    pub is_unittest_mock_assertion: bool,
     /// Set when the definition has both `*args` and `**kwargs` and both are typed
     /// `Any` — explicitly, or implicitly because they are unannotated. Per the
     /// typing spec such a signature is equivalent to `...`. Captured at definition

crates/pyrefly_types/src/lib.rs

@@ -63,3 +63,4 @@ pub mod type_var;
 pub mod type_var_tuple;
 pub mod typed_dict;
 pub mod types;
+pub mod unittest_mock;

crates/pyrefly_types/src/types.rs

@@ -75,6 +75,7 @@ use crate::type_var::Restriction;
 use crate::type_var::TypeVar;
 use crate::type_var_tuple::TypeVarTuple;
 use crate::typed_dict::TypedDict;
+use crate::unittest_mock::UnittestMock;
 
 /// An introduced synthetic variable to range over as yet unknown types.
 #[derive(Debug, Copy, Clone, Dupe, PartialEq, Eq, PartialOrd, Ord, Hash)]
@@ -961,6 +962,8 @@ pub enum Type {
     /// Instances of classes have this type, and a term of the form `C[arg1, arg2]`
     /// would have the form `Type::Type(box Type::ClassType(C, [arg1, arg2]))`.
     ClassType(ClassType),
+    /// A generated `unittest.mock` instance with a retained assertion signature.
+    UnittestMock(Box<UnittestMock>),
     /// Instances of TypedDicts have this type, and a term of the form `TD[arg1, arg2]`
     /// would have the form `Type::Type(box Type::TypedDict(TD, [arg1, arg2]))`. Note
     /// that TypedDict class definitions are still represented as `ClassDef(TD)`, just
@@ -1093,6 +1096,7 @@ impl Visit for Type {
             Type::Intersect(x) => x.visit(f),
             Type::ClassDef(x) => x.visit(f),
             Type::ClassType(x) => x.visit(f),
+            Type::UnittestMock(x) => x.visit(f),
             Type::TypedDict(x) => x.visit(f),
             Type::PartialTypedDict(x) => x.visit(f),
             Type::ShapedArray(x) => x.visit(f),
@@ -1157,6 +1161,7 @@ impl VisitMut for Type {
             Type::Intersect(x) => x.visit_mut(f),
             Type::ClassDef(x) => x.visit_mut(f),
             Type::ClassType(x) => x.visit_mut(f),
+            Type::UnittestMock(x) => x.visit_mut(f),
             Type::TypedDict(x) => x.visit_mut(f),
             Type::PartialTypedDict(x) => x.visit_mut(f),
             Type::ShapedArray(x) => x.visit_mut(f),
@@ -2470,6 +2475,7 @@ impl Type {
             | Type::Union(_)
             | Type::ClassDef(_)
             | Type::ClassType(_)
+            | Type::UnittestMock(_)
             | Type::TypedDict(_)
             | Type::PartialTypedDict(_)
             | Type::ShapedArray(_)

crates/pyrefly_types/src/unittest_mock.rs

@@ -0,0 +1,21 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+use pyrefly_derive::TypeEq;
+use pyrefly_derive::Visit;
+use pyrefly_derive::VisitMut;
+
+use crate::class::ClassType;
+use crate::types::Type;
+
+/// A generated mock retains its own behavior and a separate signature for call assertions.
+#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
+#[derive(Visit, VisitMut, TypeEq)]
+pub struct UnittestMock {
+    pub class: ClassType,
+    pub assertion_signature: Type,
+}

pyrefly/lib/alt.rs

@@ -40,4 +40,5 @@ pub mod subscript;
 pub mod traits;
 pub mod type_level_dsl;
 pub mod types;
+pub mod unittest_mock;
 pub mod unwrap;

pyrefly/lib/alt/attr.rs

@@ -635,6 +635,9 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
         context: Option<&dyn Fn() -> ErrorContext>,
         todo_ctx: &str,
     ) -> Type {
+        if let Some(signature) = self.unittest_mock_assertion_signature(base, attr_name) {
+            return signature;
+        }
         let attr_base = self.as_attribute_base(base.clone());
         let lookup_result = attr_base.clone().map_or_else(
             || LookupResult::internal_error(InternalError::AttributeBaseUndefined(base.clone())),
@@ -2662,6 +2665,7 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
 
     fn as_attribute_base1(&self, ty: Type, acc: &mut Vec<AttributeBase1>) {
         match ty {
+            Type::UnittestMock(mock) => self.as_attribute_base1(Type::ClassType(mock.class), acc),
             Type::ClassType(class_type) => {
                 if let Some(shaped_array) =
                     self.try_reproject_class_type_to_shaped_array(&class_type)

pyrefly/lib/alt/call.rs

@@ -10,7 +10,6 @@ use std::sync::Arc;
 
 use itertools::Itertools;
 use pyrefly_python::dunder;
-use pyrefly_python::module_name::ModuleName;
 use pyrefly_python::module_path::ModuleStyle;
 use pyrefly_types::dimension::ShapeError;
 use pyrefly_types::function::BodyKind;
@@ -29,7 +28,6 @@ use pyrefly_util::visit::Visit;
 use ruff_python_ast::Arguments;
 use ruff_python_ast::Expr;
 use ruff_python_ast::ExprCall;
-use ruff_python_ast::ExprStringLiteral;
 use ruff_python_ast::name::Name;
 use ruff_text_size::Ranged;
 use ruff_text_size::TextRange;
@@ -75,7 +73,6 @@ use crate::types::function::FunctionKind;
 use crate::types::keywords::KwCall;
 use crate::types::keywords::TypeMap;
 use crate::types::literal::Lit;
-use crate::types::module::ModuleType;
 use crate::types::type_var::Restriction;
 use crate::types::typed_dict::TypedDict;
 use crate::types::types::AnyStyle;
@@ -404,6 +401,9 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
 
     fn as_call_target_impl(&self, ty: Type, quantified: Option<Quantified>) -> CallTargetLookup {
         match ty {
+            Type::UnittestMock(mock) => {
+                self.as_call_target_impl(Type::ClassType(mock.class), quantified)
+            }
             Type::Callable(c) => {
                 CallTargetLookup::Ok(Box::new(CallTarget::Callable(TargetWithTParams(None, *c))))
             }
@@ -3109,7 +3109,8 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
                 )
         } else {
             self.expand_mut(&mut callee.ty);
-            self.check_unittest_mock_patch_target(&callee.ty, &callee.expr.arguments, errors);
+            let unittest_mock_patch_target =
+                self.check_unittest_mock_patch_target(&callee.ty, &callee.expr.arguments, errors);
 
             let call = CallWithTypes::new();
             let (args, kws) = if callee.ty.is_union() {
@@ -3151,6 +3152,8 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
                     errors,
                 );
                 match ty.callee_kind() {
+                _ if ty.toplevel_func_metadata().is_some_and(|metadata| metadata.flags.is_unittest_mock_assertion) =>
+                    self.call_unittest_mock_assertion(ty, &args, &kws, callee.expr, errors),
                 Some(CalleeKind::Function(FunctionKind::AssertType)) => self
                     .call_assert_type(
                         &callee.expr.arguments.args,
@@ -3428,95 +3431,24 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
                 }
             }});
             // TypeIs and TypeGuard functions return bool at runtime
-            match result {
+            let result = match result {
                 Type::TypeIs(_) | Type::TypeGuard(_) => {
                     self.heap.mk_class_type(self.stdlib.bool().clone())
                 }
                 other => other,
+            };
+            match unittest_mock_patch_target {
+                Some(target) => {
+                    self.specialize_unittest_mock_patch(result, target, &callee.expr.arguments)
+                }
+                None => result,
             }
         };
 
         let result = self.apply_django_annotate_call(result, django_annotate_call);
         self.apply_polars_call_specialization(result, polars_call)
     }
 
-    fn check_unittest_mock_patch_target(
-        &self,
-        callee_ty: &Type,
-        arguments: &Arguments,
-        errors: &ErrorCollector,
-    ) {
-        let Type::ClassType(cls) = callee_ty else {
-            return;
-        };
-        if !cls.has_qname("unittest.mock", "_patcher") {
-            return;
-        }
-        if arguments.args.len() > 1 || arguments.keywords.iter().any(|kw| kw.arg.is_none()) {
-            return;
-        }
-        if arguments
-            .find_keyword("create")
-            .is_some_and(|kw| !matches!(&kw.value, Expr::BooleanLiteral(lit) if !lit.value))
-        {
-            return;
-        }
-        let Some(target_expr) = arguments.find_argument_value("target", 0) else {
-            return;
-        };
-        let Expr::StringLiteral(ExprStringLiteral { value, .. }) = target_expr else {
-            return;
-        };
-        let range = target_expr.range();
-        let target = value.to_str();
-        let parts: Vec<&str> = target.split('.').collect();
-        if parts.len() < 2 || parts.iter().any(|p| p.is_empty()) {
-            return;
-        }
-
-        let Some((module_prefix_len, module)) = (1..parts.len()).rev().find_map(|i| {
-            let candidate = ModuleName::from_str(&parts[..i].join("."));
-            (candidate == self.module().name()
-                || self.exports.module_exists(candidate).finding().is_some())
-            .then_some((i, candidate))
-        }) else {
-            // The target may be resolved dynamically at runtime, so only check paths rooted in a
-            // module that is available in the current environment.
-            return;
-        };
-
-        let attrs = &parts[module_prefix_len..];
-        let mut base_ty = ModuleType::new_as(module).to_type(self.heap);
-        for (i, attr) in attrs.iter().enumerate() {
-            let name = Name::new(*attr);
-            // Patching a public builtin name on a module always succeeds: `mock.patch` sets
-            // `create=True` itself in that case, overriding whatever the caller passed. So
-            // `patch("mod.open")` is legal even when `mod` defines no `open` of its own.
-            // Documented at
-            // https://docs.python.org/3/library/unittest.mock.html#unittest.mock.patch ("If you
-            // are patching builtins in a module then you don't need to pass `create=True`").
-            if i + 1 == attrs.len()
-                && base_ty.as_module().is_some()
-                && !attr.starts_with('_')
-                && self.exports.export_exists(ModuleName::builtins(), &name)
-                && !self
-                    .exports
-                    .is_implicit_reexport(ModuleName::builtins(), &name)
-            {
-                return;
-            }
-            base_ty = self.type_of_attr_get(
-                &base_ty,
-                &name,
-                range,
-                errors,
-                ErrorKind::MissingAttributePatchTarget,
-                None,
-                "unittest.mock.patch target",
-            );
-        }
-    }
-
     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;
@@ -3591,7 +3523,6 @@ impl<'ctx, 'answer, Ans: LookupAnswer> AnswersSolver<'ctx, 'answer, Ans> {
             result
         })
     }
-
     pub fn freeform_call_infer(
         &self,
         ty: Type,