[red-knot] Binary operator inference: generalize code for non-instances (#17081)

## Summary

Generalize the rich-comparison fallback code for binary operator
inference. This gets rid of one `todo_type!(…)` and implements the last
remaining failing case from
https://github.com/astral-sh/ruff/issues/14200.

closes https://github.com/astral-sh/ruff/issues/14200

## Test Plan

New Markdown tests.
This commit is contained in:
David Peter 2025-03-31 13:01:25 +02:00 committed by GitHub
parent 3d1e5676fb
commit 2d7f118f52
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 59 additions and 44 deletions

View File

@ -371,6 +371,39 @@ a = NotBoolable()
10 and a and True
```
## Operations on class objects
When operating on class objects, the corresponding dunder methods are looked up on the metaclass.
```py
from __future__ import annotations
class Meta(type):
def __add__(self, other: Meta) -> int:
return 1
def __lt__(self, other: Meta) -> bool:
return True
def __getitem__(self, key: int) -> str:
return "a"
class A(metaclass=Meta): ...
class B(metaclass=Meta): ...
reveal_type(A + B) # revealed: int
# error: [unsupported-operator] "Operator `-` is unsupported between objects of type `Literal[A]` and `Literal[B]`"
reveal_type(A - B) # revealed: Unknown
reveal_type(A < B) # revealed: bool
reveal_type(A > B) # revealed: bool
# error: [unsupported-operator] "Operator `<=` is not supported for types `Literal[A]` and `Literal[B]`"
reveal_type(A <= B) # revealed: Unknown
reveal_type(A[0]) # revealed: str
```
## Unsupported
### Dunder as instance attribute

View File

@ -76,12 +76,11 @@ use crate::types::diagnostic::{
use crate::types::mro::MroErrorKind;
use crate::types::unpacker::{UnpackResult, Unpacker};
use crate::types::{
class::MetaclassErrorKind, todo_type, Class, DynamicType, FunctionType, InstanceType,
IntersectionBuilder, IntersectionType, KnownClass, KnownFunction, KnownInstanceType,
MetaclassCandidate, Parameter, ParameterForm, Parameters, SliceLiteralType, SubclassOfType,
Symbol, SymbolAndQualifiers, Truthiness, TupleType, Type, TypeAliasType, TypeAndQualifiers,
TypeArrayDisplay, TypeQualifiers, TypeVarBoundOrConstraints, TypeVarInstance, UnionBuilder,
UnionType,
class::MetaclassErrorKind, todo_type, Class, DynamicType, FunctionType, IntersectionBuilder,
IntersectionType, KnownClass, KnownFunction, KnownInstanceType, MetaclassCandidate, Parameter,
ParameterForm, Parameters, SliceLiteralType, SubclassOfType, Symbol, SymbolAndQualifiers,
Truthiness, TupleType, Type, TypeAliasType, TypeAndQualifiers, TypeArrayDisplay,
TypeQualifiers, TypeVarBoundOrConstraints, TypeVarInstance, UnionBuilder, UnionType,
};
use crate::types::{CallableType, GeneralCallableType, Signature};
use crate::unpack::{Unpack, UnpackPosition};
@ -5318,12 +5317,11 @@ impl<'db> TypeInferenceBuilder<'db> {
}
}
// Lookup the rich comparison `__dunder__` methods on instances
(Type::Instance(left_instance), Type::Instance(right_instance)) => {
let rich_comparison =
|op| self.infer_rich_comparison(left_instance, right_instance, op);
// Lookup the rich comparison `__dunder__` methods
_ => {
let rich_comparison = |op| self.infer_rich_comparison(left, right, op);
let membership_test_comparison = |op, range: TextRange| {
self.infer_membership_test_comparison(left_instance, right_instance, op, range)
self.infer_membership_test_comparison(left, right, op, range)
};
match op {
ast::CmpOp::Eq => rich_comparison(RichCompareOperator::Eq),
@ -5362,37 +5360,27 @@ impl<'db> TypeInferenceBuilder<'db> {
}
}
}
_ => match op {
ast::CmpOp::Is | ast::CmpOp::IsNot => Ok(KnownClass::Bool.to_instance(self.db())),
_ => Ok(todo_type!("Binary comparisons between more types")),
},
}
}
/// Rich comparison in Python are the operators `==`, `!=`, `<`, `<=`, `>`, and `>=`. Their
/// behaviour can be edited for classes by implementing corresponding dunder methods.
/// This function performs rich comparison between two instances and returns the resulting type.
/// This function performs rich comparison between two types and returns the resulting type.
/// see `<https://docs.python.org/3/reference/datamodel.html#object.__lt__>`
fn infer_rich_comparison(
&self,
left: InstanceType<'db>,
right: InstanceType<'db>,
left: Type<'db>,
right: Type<'db>,
op: RichCompareOperator,
) -> Result<Type<'db>, CompareUnsupportedError<'db>> {
let db = self.db();
// The following resource has details about the rich comparison algorithm:
// https://snarky.ca/unravelling-rich-comparison-operators/
let call_dunder =
|op: RichCompareOperator, left: InstanceType<'db>, right: InstanceType<'db>| {
Type::Instance(left)
.try_call_dunder(
db,
op.dunder(),
CallArgumentTypes::positional([Type::Instance(right)]),
)
.map(|outcome| outcome.return_type(db))
.ok()
};
let call_dunder = |op: RichCompareOperator, left: Type<'db>, right: Type<'db>| {
left.try_call_dunder(db, op.dunder(), CallArgumentTypes::positional([right]))
.map(|outcome| outcome.return_type(db))
.ok()
};
// The reflected dunder has priority if the right-hand side is a strict subclass of the left-hand side.
if left != right && right.is_subtype_of(db, left) {
@ -5412,8 +5400,8 @@ impl<'db> TypeInferenceBuilder<'db> {
})
.ok_or_else(|| CompareUnsupportedError {
op: op.into(),
left_ty: left.into(),
right_ty: right.into(),
left_ty: left,
right_ty: right,
})
}
@ -5423,31 +5411,25 @@ impl<'db> TypeInferenceBuilder<'db> {
/// and `<https://docs.python.org/3/reference/expressions.html#membership-test-details>`
fn infer_membership_test_comparison(
&self,
left: InstanceType<'db>,
right: InstanceType<'db>,
left: Type<'db>,
right: Type<'db>,
op: MembershipTestCompareOperator,
range: TextRange,
) -> Result<Type<'db>, CompareUnsupportedError<'db>> {
let db = self.db();
let contains_dunder = right.class().class_member(db, "__contains__").symbol;
let contains_dunder = right.class_member(db, "__contains__".into()).symbol;
let compare_result_opt = match contains_dunder {
Symbol::Type(contains_dunder, Boundness::Bound) => {
// If `__contains__` is available, it is used directly for the membership test.
contains_dunder
.try_call(
db,
CallArgumentTypes::positional([
Type::Instance(right),
Type::Instance(left),
]),
)
.try_call(db, CallArgumentTypes::positional([right, left]))
.map(|bindings| bindings.return_type(db))
.ok()
}
_ => {
// iteration-based membership test
Type::Instance(right)
right
.try_iterate(db)
.map(|_| KnownClass::Bool.to_instance(db))
.ok()
@ -5472,8 +5454,8 @@ impl<'db> TypeInferenceBuilder<'db> {
})
.ok_or_else(|| CompareUnsupportedError {
op: op.into(),
left_ty: left.into(),
right_ty: right.into(),
left_ty: left,
right_ty: right,
})
}