[ty] synthesize __replace__ for >=3.13

Author: thejchap

crates/ty_python_semantic/resources/mdtest/replace.md

@@ -0,0 +1,54 @@
+# Replace
+
+The replace function and protocol added in Python 3.13:
+<https://docs.python.org/3/whatsnew/3.13.html#copy>
+
+```toml
+[environment]
+python-version = "3.13"
+```
+
+## `replace()` function
+
+It is present in the `copy` module.
+
+```py
+from copy import replace
+```
+
+## `__replace__` protocol
+
+```toml
+[environment]
+python-version = "3.13"
+```
+
+### Dataclasses
+
+```py
+from dataclasses import dataclass
+from copy import replace
+
+@dataclass
+class Point:
+    x: int
+    y: int
+
+a = Point(1, 2)
+
+# It accepts keyword arguments
+reveal_type(a.__replace__)  # revealed: (*, x: int = int, y: int = int) -> Point
+b = a.__replace__(x=3, y=4)
+reveal_type(b)  # revealed: Point
+b = replace(a, x=3, y=4)
+reveal_type(b)  # revealed: Point
+
+# It does not require all keyword arguments
+c = a.__replace__(x=3)
+reveal_type(c)  # revealed: Point
+d = replace(a, x=3)
+reveal_type(d)  # revealed: Point
+
+e = a.__replace__(x="wrong")  # error: [invalid-argument-type]
+e = replace(a, x="wrong")  # TODO(thejchap): error: [invalid-argument-type]
+```

crates/ty_python_semantic/src/module_resolver/module.rs

@@ -264,6 +264,7 @@ pub enum KnownModule {
     Typing,
     Sys,
     Abc,
+    Copy,
     Dataclasses,
     Collections,
     Inspect,
@@ -292,6 +293,7 @@ impl KnownModule {
             Self::Sys => "sys",
             Self::Abc => "abc",
             Self::Dataclasses => "dataclasses",
+            Self::Copy => "copy",
             Self::Collections => "collections",
             Self::Inspect => "inspect",
             Self::TypeCheckerInternals => "_typeshed._type_checker_internals",

crates/ty_python_semantic/src/types/call/bind.rs

@@ -12,9 +12,10 @@ use ruff_db::parsed::parsed_module;
 use smallvec::{SmallVec, smallvec, smallvec_inline};
 
 use super::{Argument, CallArguments, CallError, CallErrorKind, InferContext, Signature, Type};
+use crate::Program;
 use crate::db::Db;
 use crate::dunder_all::dunder_all_names;
-use crate::place::{Boundness, Place};
+use crate::place::{Boundness, Place, PlaceAndQualifiers};
 use crate::types::diagnostic::{
     CALL_NON_CALLABLE, CONFLICTING_ARGUMENT_FORMS, INVALID_ARGUMENT_TYPE, MISSING_ARGUMENT,
     NO_MATCHING_OVERLOAD, PARAMETER_ALREADY_ASSIGNED, TOO_MANY_POSITIONAL_ARGUMENTS,
@@ -29,11 +30,11 @@ use crate::types::signatures::{Parameter, ParameterForm, Parameters};
 use crate::types::tuple::{Tuple, TupleLength, TupleType};
 use crate::types::{
     BoundMethodType, ClassLiteral, DataclassParams, FieldInstance, KnownClass, KnownInstanceType,
-    MethodWrapperKind, PropertyInstanceType, SpecialFormType, TypeMapping, UnionType,
-    WrapperDescriptorKind, enums, ide_support, todo_type,
+    MemberLookupPolicy, MethodWrapperKind, PropertyInstanceType, SpecialFormType, TypeMapping,
+    UnionType, WrapperDescriptorKind, enums, ide_support, todo_type,
 };
 use ruff_db::diagnostic::{Annotation, Diagnostic, SubDiagnostic, SubDiagnosticSeverity};
-use ruff_python_ast as ast;
+use ruff_python_ast::{self as ast, PythonVersion};
 
 /// Binding information for a possible union of callables. At a call site, the arguments must be
 /// compatible with _all_ of the types in the union for the call to be valid.
@@ -948,6 +949,37 @@ impl<'db> Bindings<'db> {
                             }
                         }
 
+                        Some(KnownFunction::Replace)
+                            if Program::get(db).python_version(db) >= PythonVersion::PY313 =>
+                        {
+                            let [Some(instance_ty), _] = overload.parameter_types() else {
+                                continue;
+                            };
+                            let PlaceAndQualifiers {
+                                place: Place::Type(Type::Callable(attr_ty), _),
+                                qualifiers: _,
+                            } = instance_ty.class_member_with_policy(
+                                db,
+                                "__replace__".into(),
+                                MemberLookupPolicy::MRO_NO_OBJECT_FALLBACK,
+                            )
+                            else {
+                                continue;
+                            };
+                            let Some(
+                                _signature @ Signature {
+                                    return_ty: Some(return_ty),
+                                    ..
+                                },
+                            ) = attr_ty.signatures(db).overloads.first()
+                            else {
+                                continue;
+                            };
+                            // TODO(thejchap): set param types on overload
+                            // println!("params: {:?}", signature.parameters());
+                            overload.set_return_type(*return_ty);
+                        }
+
                         _ => {
                             // Ideally, either the implementation, or exactly one of the overloads
                             // of the function can have the dataclass_transform decorator applied.

crates/ty_python_semantic/src/types/class.rs

@@ -1596,7 +1596,7 @@ impl<'db> ClassLiteral<'db> {
 
         let field_policy = CodeGeneratorKind::from_class(db, self)?;
 
-        let signature_from_fields = |mut parameters: Vec<_>| {
+        let signature_from_fields = |mut parameters: Vec<_>, return_ty: Option<Type<'db>>| {
             let mut kw_only_field_seen = false;
             for (
                 field_name,
@@ -1669,21 +1669,27 @@ impl<'db> ClassLiteral<'db> {
                     }
                 }
 
-                let mut parameter = if kw_only_field_seen {
+                // For the `__replace__` signature, force to kw only
+                let mut parameter = if kw_only_field_seen || name == "__replace__" {
                     Parameter::keyword_only(field_name)
                 } else {
                     Parameter::positional_or_keyword(field_name)
                 }
                 .with_annotated_type(field_ty);
 
-                if let Some(default_ty) = default_ty {
+                if name == "__replace__" {
+                    // When replacing, we know there is a default value for the field
+                    // (the value that is currently assigned to the field)
+                    // assume this to be the declared type of the field
+                    parameter = parameter.with_default_type(field_ty);
+                } else if let Some(default_ty) = default_ty {
                     parameter = parameter.with_default_type(default_ty);
                 }
 
                 parameters.push(parameter);
             }
 
-            let mut signature = Signature::new(Parameters::new(parameters), Some(Type::none(db)));
+            let mut signature = Signature::new(Parameters::new(parameters), return_ty);
             signature.inherited_generic_context = self.generic_context(db);
             Some(CallableType::function_like(db, signature))
         };
@@ -1705,12 +1711,12 @@ impl<'db> ClassLiteral<'db> {
                         db,
                         self.apply_optional_specialization(db, specialization),
                     ));
-                signature_from_fields(vec![self_parameter])
+                signature_from_fields(vec![self_parameter], Some(Type::none(db)))
             }
             (CodeGeneratorKind::NamedTuple, "__new__") => {
                 let cls_parameter = Parameter::positional_or_keyword(Name::new_static("cls"))
                     .with_annotated_type(KnownClass::Type.to_instance(db));
-                signature_from_fields(vec![cls_parameter])
+                signature_from_fields(vec![cls_parameter], Some(Type::none(db)))
             }
             (CodeGeneratorKind::DataclassLike, "__lt__" | "__le__" | "__gt__" | "__ge__") => {
                 if !has_dataclass_param(DataclassParams::ORDER) {
@@ -1745,6 +1751,18 @@ impl<'db> ClassLiteral<'db> {
                     .place
                     .ignore_possibly_unbound()
             }
+            (CodeGeneratorKind::DataclassLike, "__replace__")
+                if Program::get(db).python_version(db) >= PythonVersion::PY313 =>
+            {
+                let self_parameter =
+                    Parameter::positional_or_keyword(Name::new_static("self")).with_annotated_type(
+                        Type::instance(db, self.apply_optional_specialization(db, specialization)),
+                    );
+                let instance_ty =
+                    Type::instance(db, self.apply_optional_specialization(db, specialization));
+
+                signature_from_fields(vec![self_parameter], Some(instance_ty))
+            }
             (CodeGeneratorKind::DataclassLike, "__setattr__") => {
                 if has_dataclass_param(DataclassParams::FROZEN) {
                     let signature = Signature::new(

crates/ty_python_semantic/src/types/function.rs

@@ -1053,6 +1053,9 @@ pub enum KnownFunction {
     /// `dataclasses.field`
     Field,
 
+    /// `copy.replace`
+    Replace,
+
     /// `inspect.getattr_static`
     GetattrStatic,
 
@@ -1133,6 +1136,9 @@ impl KnownFunction {
             Self::Dataclass | Self::Field => {
                 matches!(module, KnownModule::Dataclasses)
             }
+            Self::Replace => {
+                matches!(module, KnownModule::Copy)
+            }
             Self::GetattrStatic => module.is_inspect(),
             Self::IsAssignableTo
             | Self::IsDisjointFrom
@@ -1412,6 +1418,7 @@ pub(crate) mod tests {
                 KnownFunction::AbstractMethod => KnownModule::Abc,
 
                 KnownFunction::Dataclass | KnownFunction::Field => KnownModule::Dataclasses,
+                KnownFunction::Replace => KnownModule::Copy,
 
                 KnownFunction::GetattrStatic => KnownModule::Inspect,