fix Feature suggestion: special case _replace in NamedTuple #3960

pyrefly/lib/alt/class/named_tuple.rs

@@ -27,6 +27,8 @@ use crate::types::class::ClassType;
 use crate::types::literal::Lit;
 use crate::types::types::Type;
 
+const NAMED_TUPLE_REPLACE: Name = Name::new_static("_replace");
+
 impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
     pub fn get_named_tuple_elements(&self, cls: &Class, errors: &ErrorCollector) -> SmallSet<Name> {
         let Some(class_fields) = self.get_class_fields(cls) else {
@@ -145,6 +147,28 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
         ClassSynthesizedField::new(ty)
     }
 
+    fn get_named_tuple_replace(
+        &self,
+        cls: &Class,
+        elements: &SmallSet<Name>,
+        has_dynamic_fields: bool,
+    ) -> ClassSynthesizedField {
+        let mut params = vec![self.class_self_param(cls, true)];
+        if has_dynamic_fields {
+            params.push(Param::Kwargs(None, self.heap.mk_any_implicit()));
+        } else {
+            params.extend(elements.iter().map(|name| {
+                let ty = match self.get_non_synthesized_class_member(cls, name) {
+                    None => self.heap.mk_any_implicit(),
+                    Some(c) => c.as_named_tuple_type(),
+                };
+                Param::KwOnly(name.clone(), ty, Required::Optional(None))
+            }));
+        }
+        let ty = self.synthesized_method(cls, NAMED_TUPLE_REPLACE, params, self.instantiate(cls));
+        ClassSynthesizedField::new(ty)
+    }
+
     fn get_named_tuple_iter(
         &self,
         cls: &Class,
@@ -190,6 +214,7 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
         Some(ClassSynthesizedFields::new(smallmap! {
             dunder::NEW => self.get_named_tuple_new(cls, &named_tuple.elements, named_tuple.has_dynamic_fields),
             dunder::INIT => self.get_named_tuple_init(cls),
+            NAMED_TUPLE_REPLACE => self.get_named_tuple_replace(cls, &named_tuple.elements, named_tuple.has_dynamic_fields),
             dunder::MATCH_ARGS => self.get_named_tuple_match_args(&named_tuple.elements),
             dunder::ITER => self.get_named_tuple_iter(cls, &named_tuple.elements)
         }))

pyrefly/lib/test/named_tuple.rs

@@ -102,6 +102,40 @@ Point3(1)  # E: Missing argument `y` in function `Point3.__new__`
     "#,
 );
 
+testcase!(
+    test_named_tuple_replace,
+    r#"
+from typing import NamedTuple, assert_type
+
+class Point(NamedTuple):
+    x: int
+    y: str
+
+p = Point(1, "")
+assert_type(p._replace(x=2), Point)
+p._replace()
+p._replace(z=5)  # E: Unexpected keyword argument `z`
+p._replace(x="str")  # E: is not assignable to parameter `x`
+p._replace(1)  # E: Expected 0 positional arguments
+    "#,
+);
+
+testcase!(
+    test_named_tuple_functional_replace,
+    r#"
+from collections import namedtuple
+from typing import NamedTuple
+
+Point1 = namedtuple("Point1", ["x", "y"])
+Point2 = NamedTuple("Point2", [("x", int), ("y", str)])
+
+Point1(1, "")._replace(x="anything")
+Point1(1, "")._replace(z=5)  # E: Unexpected keyword argument `z`
+Point2(1, "")._replace(x="str")  # E: is not assignable to parameter `x`
+Point2(1, "")._replace(y="str")
+    "#,
+);
+
 // Regression test for https://github.com/facebook/pyrefly/issues/2811
 testcase!(
     test_inline_collections_namedtuple_constructor,