[red-knot] Dataclasses: support order=True

Author: sharkdp

.github/workflows/mypy_primer.yaml

@@ -68,7 +68,7 @@ jobs:
             --type-checker knot \
             --old base_commit \
             --new "$GITHUB_SHA" \
-            --project-selector '/(mypy_primer|black|pyp|git-revise|zipp|arrow|isort|itsdangerous|rich|packaging|pybind11|pyinstrument|typeshed-stats|scrapy|werkzeug|bidict|async-utils)$' \
+            --project-selector '/(mypy_primer|black|pyp|git-revise|zipp|arrow|isort|itsdangerous|rich|packaging|pybind11|pyinstrument|typeshed-stats|scrapy|werkzeug|bidict|async-utils|python-chess|dacite|python-htmlgen|paroxython|porcupine|psycopg)$' \
             --output concise \
             --debug > mypy_primer.diff || [ $? -eq 1 ]
 

crates/red_knot_python_semantic/resources/mdtest/dataclasses.md

@@ -91,6 +91,112 @@ repr(C())
 C() == C()
 ```
 
+## Other dataclass parameters
+
+### `repr`
+
+A custom `__repr__` method is generated by default. It can be disabled by passing `repr=False`, but
+in that case `__repr__` is still available via `object.__repr__`:
+
+```py
+from dataclasses import dataclass
+
+@dataclass(repr=False)
+class WithoutRepr:
+    x: int
+
+reveal_type(WithoutRepr(1).__repr__)  # revealed: bound method WithoutRepr.__repr__() -> str
+```
+
+### `eq`
+
+The same is true for `__eq__`. Setting `eq=False` disables the generated `__eq__` method, but
+`__eq__` is still available via `object.__eq__`:
+
+```py
+from dataclasses import dataclass
+
+@dataclass(eq=False)
+class WithoutEq:
+    x: int
+
+reveal_type(WithoutEq(1) == WithoutEq(2))  # revealed: bool
+```
+
+### `order`
+
+`order` is set to `False` by default. If `order=True`, `__lt__`, `__le__`, `__gt__`, and `__ge__`
+methods will be generated:
+
+```py
+from dataclasses import dataclass
+
+@dataclass
+class WithoutOrder:
+    x: int
+
+WithoutOrder(1) < WithoutOrder(2)  # error: [unsupported-operator]
+WithoutOrder(1) <= WithoutOrder(2)  # error: [unsupported-operator]
+WithoutOrder(1) > WithoutOrder(2)  # error: [unsupported-operator]
+WithoutOrder(1) >= WithoutOrder(2)  # error: [unsupported-operator]
+
+@dataclass(order=True)
+class WithOrder:
+    x: int
+
+WithOrder(1) < WithOrder(2)
+WithOrder(1) <= WithOrder(2)
+WithOrder(1) > WithOrder(2)
+WithOrder(1) >= WithOrder(2)
+```
+
+Comparisons are only allowed for `WithOrder` instances:
+
+```py
+WithOrder(1) < 2  # error: [unsupported-operator]
+WithOrder(1) <= 2  # error: [unsupported-operator]
+WithOrder(1) > 2  # error: [unsupported-operator]
+WithOrder(1) >= 2  # error: [unsupported-operator]
+```
+
+This also works for generic dataclasses:
+
+```py
+from dataclasses import dataclass
+
+@dataclass(order=True)
+class GenericWithOrder[T]:
+    x: T
+
+GenericWithOrder[int](1) < GenericWithOrder[int](1)
+
+GenericWithOrder[int](1) < GenericWithOrder[str]("a")  # error: [unsupported-operator]
+```
+
+### `unsafe_hash`
+
+To do
+
+### `frozen`
+
+To do
+
+### `match_args`
+
+To do
+
+### `kw_only`
+
+To do
+
+### `slots`
+
+To do
+
+### `weakref_slot`
+
+To do
+
 ## Inheritance
 
 ### Normal class inheriting from a dataclass
@@ -168,13 +274,30 @@ reveal_type(d_int.description)  # revealed: str
 DataWithDescription[int](None, "description")
 ```
 
-## Frozen instances
+## Descriptor-typed fields
 
-To do
+```py
+from dataclasses import dataclass
 
-## Descriptor-typed fields
+class Descriptor:
+    _value: int = 0
 
-To do
+    def __get__(self, instance, owner) -> str:
+        return str(self._value)
+
+    def __set__(self, instance, value: int) -> None:
+        self._value = value
+
+@dataclass
+class C:
+    d: Descriptor = Descriptor()
+
+c = C(1)
+reveal_type(c.d)  # revealed: str
+
+# TODO: should be an error
+C("a")
+```
 
 ## `dataclasses.field`
 
@@ -197,18 +320,61 @@ class C:
 reveal_type(C.__init__)  # revealed: (*args: Any, **kwargs: Any) -> None
 ```
 
-### Dataclass with `init=False`
+### Dataclass with custom `__init__` method
 
-To do
+If a class already defines `__init__`, it is not replaced by the `dataclass` decorator.
 
-### Dataclass with custom `__init__` method
+```py
+from dataclasses import dataclass
 
-To do
+@dataclass(init=True)
+class C:
+    x: str
+
+    def __init__(self, x: int) -> None:
+        self.x = str(x)
+
+C(1)  # OK
+
+# TODO: should be an error
+C("a")
+```
+
+Similarly, if we set `init=False`, we still recognize the custom `__init__` method:
+
+```py
+@dataclass(init=False)
+class D:
+    def __init__(self, x: int) -> None:
+        self.x = str(x)
+
+D(1)  # OK
+D()  # error: [missing-argument]
+```
 
 ### Dataclass with `ClassVar`s
 
 To do
 
+### Return type of `dataclass(...)`
+
+A call like `dataclass(order=True)` returns a callable itself, which is then used as the decorator.
+We can store the callable in a variable and later use it as a decorator:
+
+```py
+from dataclasses import dataclass
+
+dataclass_with_order = dataclass(order=True)
+
+reveal_type(dataclass_with_order)  # revealed: <decorator produced by dataclasses.dataclass>
+
+@dataclass_with_order
+class C:
+    x: int
+
+C(1) < C(2)  # ok
+```
+
 ### Using `dataclass` as a function
 
 To do

crates/red_knot_python_semantic/src/types/class.rs

@@ -297,10 +297,15 @@ impl<'db> ClassType<'db> {
     /// Returns [`Symbol::Unbound`] if `name` cannot be found in this class's scope
     /// directly. Use [`ClassType::class_member`] if you require a method that will
     /// traverse through the MRO until it finds the member.
-    pub(super) fn own_class_member(self, db: &'db dyn Db, name: &str) -> SymbolAndQualifiers<'db> {
+    pub(super) fn own_class_member(
+        self,
+        db: &'db dyn Db,
+        specialization: Option<Specialization<'db>>,
+        name: &str,
+    ) -> SymbolAndQualifiers<'db> {
         let (class_literal, _) = self.class_literal(db);
         class_literal
-            .own_class_member(db, name)
+            .own_class_member(db, specialization, name)
             .map_type(|ty| self.specialize_type(db, ty))
     }
 
@@ -745,7 +750,8 @@ impl<'db> ClassLiteralType<'db> {
                     }
 
                     lookup_result = lookup_result.or_else(|lookup_error| {
-                        lookup_error.or_fall_back_to(db, class.own_class_member(db, name))
+                        lookup_error
+                            .or_fall_back_to(db, class.own_class_member(db, specialization, name))
                     });
                 }
             }
@@ -790,23 +796,38 @@ impl<'db> ClassLiteralType<'db> {
     /// Returns [`Symbol::Unbound`] if `name` cannot be found in this class's scope
     /// directly. Use [`ClassLiteralType::class_member`] if you require a method that will
     /// traverse through the MRO until it finds the member.
-    pub(super) fn own_class_member(self, db: &'db dyn Db, name: &str) -> SymbolAndQualifiers<'db> {
+    pub(super) fn own_class_member(
+        self,
+        db: &'db dyn Db,
+        specialization: Option<Specialization<'db>>,
+        name: &str,
+    ) -> SymbolAndQualifiers<'db> {
         if let Some(metadata) = self.dataclass_metadata(db) {
-            if name == "__init__" {
-                if 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)),
+            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)),
+                );
+
+                return Symbol::bound(Type::Callable(CallableType::new(db, init_signature))).into();
+            } else if matches!(name, "__lt__" | "__le__" | "__gt__" | "__ge__") {
+                if metadata.contains(DataclassMetadata::ORDER) {
+                    let signature = Signature::new(
+                        Parameters::new([Parameter::positional_or_keyword(Name::new_static(
+                            "other",
+                        ))
+                        .with_annotated_type(Type::instance(
+                            self.apply_optional_specialization(db, specialization),
+                        ))]),
+                        Some(KnownClass::Bool.to_instance(db)),
                     );
-
-                    return Symbol::bound(Type::Callable(CallableType::new(db, init_signature)))
-                        .into();
+                    return Symbol::bound(Type::Callable(CallableType::new(db, signature))).into();
                 }
             }
         }

crates/red_knot_python_semantic/src/types/generics.rs

@@ -9,7 +9,7 @@ use crate::types::{
 use crate::Db;
 
 /// A list of formal type variables for a generic function, class, or type alias.
-#[salsa::tracked(debug)]
+#[salsa::interned(debug)]
 pub struct GenericContext<'db> {
     #[return_ref]
     pub(crate) variables: Box<[TypeVarInstance<'db>]>,
@@ -21,7 +21,7 @@ impl<'db> GenericContext<'db> {
         index: &'db SemanticIndex<'db>,
         type_params_node: &ast::TypeParams,
     ) -> Self {
-        let variables = type_params_node
+        let variables: Box<[_]> = type_params_node
             .iter()
             .filter_map(|type_param| Self::variable_from_type_param(db, index, type_param))
             .collect();
@@ -100,7 +100,7 @@ impl<'db> GenericContext<'db> {
 }
 
 /// An assignment of a specific type to each type variable in a generic scope.
-#[salsa::tracked(debug)]
+#[salsa::interned(debug)]
 pub struct Specialization<'db> {
     pub(crate) generic_context: GenericContext<'db>,
     #[return_ref]
@@ -122,7 +122,7 @@ impl<'db> Specialization<'db> {
     /// That lets us produce the generic alias `A[int]`, which is the corresponding entry in the
     /// MRO of `B[int]`.
     pub(crate) fn apply_specialization(self, db: &'db dyn Db, other: Specialization<'db>) -> Self {
-        let types = self
+        let types: Box<[_]> = self
             .types(db)
             .into_iter()
             .map(|ty| ty.apply_specialization(db, other))
@@ -131,7 +131,7 @@ impl<'db> Specialization<'db> {
     }
 
     pub(crate) fn normalized(self, db: &'db dyn Db) -> Self {
-        let types = self.types(db).iter().map(|ty| ty.normalized(db)).collect();
+        let types: Box<[_]> = self.types(db).iter().map(|ty| ty.normalized(db)).collect();
         Self::new(db, self.generic_context(db), types)
     }
 

crates/red_knot_python_semantic/src/types/slots.rs

@@ -24,7 +24,8 @@ enum SlotsKind {
 
 impl SlotsKind {
     fn from(db: &dyn Db, base: ClassLiteralType) -> Self {
-        let Symbol::Type(slots_ty, bound) = base.own_class_member(db, "__slots__").symbol else {
+        let Symbol::Type(slots_ty, bound) = base.own_class_member(db, None, "__slots__").symbol
+        else {
             return Self::NotSpecified;
         };