[ty] Detect overloads decorated with @dataclass_transform

Author: abhijeetbodas2001

crates/ty_python_semantic/resources/mdtest/dataclass_transform.md

@@ -283,11 +283,11 @@ class D2:
 
 # TODO: these should not be errors
 D1("a")  # error: [too-many-positional-arguments]
-D2("a")  # error: [too-many-positional-arguments]
+D2("a")
 
 # TODO: these should be invalid-argument-type errors
 D1(1.2)  # error: [too-many-positional-arguments]
-D2(1.2)  # error: [too-many-positional-arguments]
+D2(1.2)  # error: [invalid-argument-type]
 ```
 
 [`typing.dataclass_transform`]: https://docs.python.org/3/library/typing.html#typing.dataclass_transform

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

@@ -20,8 +20,8 @@ use crate::types::generics::{Specialization, SpecializationBuilder, Specializati
 use crate::types::signatures::{Parameter, ParameterForm};
 use crate::types::{
     todo_type, BoundMethodType, DataclassParams, DataclassTransformerParams, FunctionDecorators,
-    KnownClass, KnownFunction, KnownInstanceType, MethodWrapperKind, PropertyInstanceType,
-    TupleType, UnionType, WrapperDescriptorKind,
+    FunctionType, KnownClass, KnownFunction, KnownInstanceType, MethodWrapperKind,
+    PropertyInstanceType, TupleType, UnionType, WrapperDescriptorKind,
 };
 use ruff_db::diagnostic::{Annotation, Severity, SubDiagnostic};
 use ruff_python_ast as ast;
@@ -770,29 +770,50 @@ impl<'db> Bindings<'db> {
                         }
 
                         _ => {
-                            if let Some(params) = function_type.dataclass_transformer_params(db) {
-                                // This is a call to a custom function that was decorated with `@dataclass_transformer`.
-                                // If this function was called with a keyword argument like `order=False`, we extract
-                                // the argument type and overwrite the corresponding flag in `dataclass_params` after
-                                // constructing them from the `dataclass_transformer`-parameter defaults.
-
-                                let mut dataclass_params = DataclassParams::from(params);
+                            let mut handle_dataclass_transformer_params =
+                                |function_type: &FunctionType| {
+                                    if let Some(params) =
+                                        function_type.dataclass_transformer_params(db)
+                                    {
+                                        // This is a call to a custom function that was decorated with `@dataclass_transformer`.
+                                        // If this function was called with a keyword argument like `order=False`, we extract
+                                        // the argument type and overwrite the corresponding flag in `dataclass_params` after
+                                        // constructing them from the `dataclass_transformer`-parameter defaults.
+
+                                        let mut dataclass_params = DataclassParams::from(params);
+
+                                        if let Some(Some(Type::BooleanLiteral(order))) =
+                                            callable_signature.iter().nth(overload_index).and_then(
+                                                |signature| {
+                                                    let (idx, _) = signature
+                                                        .parameters()
+                                                        .keyword_by_name("order")?;
+                                                    overload.parameter_types().get(idx)
+                                                },
+                                            )
+                                        {
+                                            dataclass_params.set(DataclassParams::ORDER, *order);
+                                        }
 
-                                if let Some(Some(Type::BooleanLiteral(order))) = callable_signature
+                                        overload.set_return_type(Type::DataclassDecorator(
+                                            dataclass_params,
+                                        ));
+                                    }
+                                };
+
+                            // Ideally, either the implementation, or exactly one of the overloads
+                            // of the function can have the dataclass_transform decorator applied.
+                            // However, we do not yet enforce this, and in the case of multiple
+                            // applications of the decorator, we will only consider the last one
+                            // for the return value, since the prior ones will be over-written.
+                            if let Some(overloaded) = function_type.to_overloaded(db) {
+                                overloaded
+                                    .overloads
                                     .iter()
-                                    .nth(overload_index)
-                                    .and_then(|signature| {
-                                        let (idx, _) =
-                                            signature.parameters().keyword_by_name("order")?;
-                                        overload.parameter_types().get(idx)
-                                    })
-                                {
-                                    dataclass_params.set(DataclassParams::ORDER, *order);
-                                }
+                                    .for_each(&mut handle_dataclass_transformer_params);
+                            };
 
-                                overload
-                                    .set_return_type(Type::DataclassDecorator(dataclass_params));
-                            }
+                            handle_dataclass_transformer_params(&function_type);
                         }
                     },