Initial version

Author: sharkdp

crates/red_knot_python_semantic/resources/mdtest/dataclasses.md

@@ -40,19 +40,79 @@ The signature of the `__init__` method is generated based on the classes attribu
 calls are not valid:
 
 ```py
-# TODO: should be an error: too few arguments
+# error: [missing-argument]
 Person()
 
-# TODO: should be an error: too many arguments
+# error: [too-many-positional-arguments]
 Person("Eve", 20, "too many arguments")
 
-# TODO: should be an error: wrong argument type
+# error: [invalid-argument-type]
 Person("Eve", "string instead of int")
 
-# TODO: should be an error: wrong argument types
+# error: [invalid-argument-type]
+# error: [invalid-argument-type]
 Person(20, "Eve")
 ```
 
+## Signature of `__init__`
+
+TODO: All of the following tests are missing the `self` argument in the `__init__` signature.
+
+Declarations in the class body are used to generate the signature of the `__init__` method. If the
+attributes are not just declarations, but also bindings, the type inferred from bindings is used as
+the default value.
+
+```py
+from dataclasses import dataclass
+
+@dataclass
+class D1:
+    x: int
+    y: str = "default"
+    z: int | None = 1 + 2
+
+reveal_type(D1.__init__)  # revealed: (x: int, y: str = Literal["default"], z: int | None = Literal[3]) -> None
+```
+
+This also works if the declaration and binding are split:
+
+```py
+@dataclass
+class D2:
+    x: int | None
+    x = None
+
+reveal_type(D2.__init__)  # revealed: (x: int | None = None) -> None
+```
+
+Function declarations do not affect the signature of `__init__`:
+
+```py
+@dataclass
+class D3:
+    x: int
+
+    def y(self) -> str:
+        return ""
+
+reveal_type(D3.__init__)  # revealed: (x: int) -> None
+```
+
+Implicit instance attributes also do not affect the signature of `__init__`:
+
+```py
+@dataclass
+class D4:
+    x: int
+
+    def f(self, y: str) -> None:
+        self.y: str = y
+
+reveal_type(D4(1).y)  # revealed: str
+
+reveal_type(D4.__init__)  # revealed: (x: int) -> None
+```
+
 ## `@dataclass` calls with arguments
 
 The `@dataclass` decorator can take several arguments to customize the existence of the generated
@@ -228,7 +288,8 @@ class Derived(Base):
 
 d = Derived("a")
 
-# TODO: should be an error:
+# error: [too-many-positional-arguments]
+# error: [invalid-argument-type]
 Derived(1, "a")
 ```
 
@@ -245,6 +306,9 @@ class Base:
 class Derived(Base):
     y: str
 
+# TODO: no errors
+# error: [too-many-positional-arguments]
+# error: [invalid-argument-type]
 d = Derived(1, "a")  # OK
 
 reveal_type(d.x)  # revealed: int
@@ -292,10 +356,12 @@ class Descriptor:
 class C:
     d: Descriptor = Descriptor()
 
+# TODO: no error
+# error: [invalid-argument-type]
 c = C(1)
 reveal_type(c.d)  # revealed: str
 
-# TODO: should be an error
+# error: [invalid-argument-type]
 C("a")
 ```
 
@@ -316,8 +382,7 @@ import dataclasses
 class C:
     x: str
 
-# TODO: should show the proper signature
-reveal_type(C.__init__)  # revealed: (*args: Any, **kwargs: Any) -> None
+reveal_type(C.__init__)  # revealed: (x: str) -> None
 ```
 
 ### Dataclass with custom `__init__` method
@@ -334,6 +399,8 @@ class C:
     def __init__(self, x: int) -> None:
         self.x = str(x)
 
+# TODO: no error
+# error: [invalid-argument-type]
 C(1)  # OK
 
 # TODO: should be an error
@@ -399,8 +466,8 @@ reveal_type(Person.__mro__)  # revealed: tuple[Literal[Person], Literal[object]]
 The generated methods have the following signatures:
 
 ```py
-# TODO: proper signature
-reveal_type(Person.__init__)  # revealed: (*args: Any, **kwargs: Any) -> None
+# TODO: `self` is missing here
+reveal_type(Person.__init__)  # revealed: (name: str, age: int | None = None) -> None
 
 reveal_type(Person.__repr__)  # revealed: def __repr__(self) -> str
 

crates/red_knot_python_semantic/src/semantic_index/use_def.rs

@@ -429,6 +429,14 @@ impl<'db> UseDefMap<'db> {
         self.declarations_iterator(declarations)
     }
 
+    pub(crate) fn all_public_declarations<'map>(
+        &'map self,
+    ) -> impl Iterator<Item = (ScopedSymbolId, DeclarationsIterator<'map, 'db>)> + 'map {
+        (0..self.public_symbols.len())
+            .map(ScopedSymbolId::from_usize)
+            .map(|symbol_id| (symbol_id, self.public_declarations(symbol_id)))
+    }
+
     /// This function is intended to be called only once inside `TypeInferenceBuilder::infer_function_body`.
     pub(crate) fn can_implicit_return(&self, db: &dyn crate::Db) -> bool {
         !self

crates/red_knot_python_semantic/src/types/class.rs

@@ -804,16 +804,22 @@ impl<'db> ClassLiteralType<'db> {
     ) -> SymbolAndQualifiers<'db> {
         if let Some(metadata) = self.dataclass_metadata(db) {
             if name == "__init__" && metadata.contains(DataclassMetadata::INIT) {
-                // TODO: Generate the signature from the attributes on the class
-                let init_signature = Signature::new(
-                    Parameters::new([
-                        Parameter::variadic(Name::new_static("args"))
-                            .with_annotated_type(Type::any()),
-                        Parameter::keyword_variadic(Name::new_static("kwargs"))
-                            .with_annotated_type(Type::any()),
-                    ]),
-                    Some(Type::none(db)),
-                );
+                // dbg!(&attributes);
+                let mut parameters = vec![];
+
+                for (name, attr_ty, default_ty) in self.own_class_body_instance_attributes(db) {
+                    let mut parameter =
+                        Parameter::positional_or_keyword(name).with_annotated_type(attr_ty);
+
+                    if let Some(default_ty) = default_ty {
+                        parameter = parameter.with_default_type(default_ty);
+                    }
+
+                    parameters.push(parameter);
+                }
+
+                let init_signature =
+                    Signature::new(Parameters::new(parameters), Some(Type::none(db)));
 
                 return Symbol::bound(Type::Callable(CallableType::new(db, init_signature))).into();
             } else if matches!(name, "__lt__" | "__le__" | "__gt__" | "__ge__") {
@@ -836,6 +842,40 @@ impl<'db> ClassLiteralType<'db> {
         class_symbol(db, body_scope, name)
     }
 
+    fn own_class_body_instance_attributes(
+        self,
+        db: &'db dyn Db,
+    ) -> Vec<(Name, Type<'db>, Option<Type<'db>>)> {
+        let mut attributes = vec![]; // TODO: iterator
+
+        let class_body_scope = self.body_scope(db);
+        let table = symbol_table(db, class_body_scope);
+
+        let use_def = use_def_map(db, class_body_scope);
+        for (symbol_id, declarations) in use_def.all_public_declarations() {
+            let symbol = table.symbol(symbol_id);
+
+            if let Ok(attr) = symbol_from_declarations(db, declarations) {
+                if attr.is_class_var() {
+                    continue;
+                }
+
+                if let Some(attr_ty) = attr.symbol.ignore_possibly_unbound() {
+                    if attr_ty.is_function_literal() {
+                        continue;
+                    }
+
+                    let bindings = use_def.public_bindings(symbol_id);
+                    let default_ty = symbol_from_bindings(db, bindings).ignore_possibly_unbound();
+
+                    attributes.push((symbol.name().clone(), attr_ty, default_ty));
+                }
+            }
+        }
+
+        attributes
+    }
+
     /// Returns the `name` attribute of an instance of this class.
     ///
     /// The attribute could be defined in the class body, but it could also be an implicitly