[red-knot] feat: Inference for `BytesLiteral` comparisons (#13746)

Implements inference for `BytesLiteral` comparisons along the lines of
https://github.com/astral-sh/ruff/pull/13634.

closes #13687

Co-authored-by: Alex Waygood <Alex.Waygood@Gmail.com>
This commit is contained in:
David Peter 2024-10-14 14:01:23 +02:00 committed by GitHub
parent 9bb4722ebf
commit 93097f1c53
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 88 additions and 0 deletions

View File

@ -0,0 +1,43 @@
### Comparison: Byte literals
These tests assert that we infer precise `Literal` types for comparisons between objects
inferred as having `Literal` bytes types:
```py
reveal_type(b"abc" == b"abc") # revealed: Literal[True]
reveal_type(b"abc" == b"ab") # revealed: Literal[False]
reveal_type(b"abc" != b"abc") # revealed: Literal[False]
reveal_type(b"abc" != b"ab") # revealed: Literal[True]
reveal_type(b"abc" < b"abd") # revealed: Literal[True]
reveal_type(b"abc" < b"abb") # revealed: Literal[False]
reveal_type(b"abc" <= b"abc") # revealed: Literal[True]
reveal_type(b"abc" <= b"abb") # revealed: Literal[False]
reveal_type(b"abc" > b"abd") # revealed: Literal[False]
reveal_type(b"abc" > b"abb") # revealed: Literal[True]
reveal_type(b"abc" >= b"abc") # revealed: Literal[True]
reveal_type(b"abc" >= b"abd") # revealed: Literal[False]
reveal_type(b"" in b"") # revealed: Literal[True]
reveal_type(b"" in b"abc") # revealed: Literal[True]
reveal_type(b"abc" in b"") # revealed: Literal[False]
reveal_type(b"ab" in b"abc") # revealed: Literal[True]
reveal_type(b"abc" in b"abc") # revealed: Literal[True]
reveal_type(b"d" in b"abc") # revealed: Literal[False]
reveal_type(b"ac" in b"abc") # revealed: Literal[False]
reveal_type(b"\x81\x82" in b"\x80\x81\x82") # revealed: Literal[True]
reveal_type(b"\x82\x83" in b"\x80\x81\x82") # revealed: Literal[False]
reveal_type(b"ab" not in b"abc") # revealed: Literal[False]
reveal_type(b"ac" not in b"abc") # revealed: Literal[True]
reveal_type(b"abc" is b"abc") # revealed: bool
reveal_type(b"abc" is b"ab") # revealed: Literal[False]
reveal_type(b"abc" is not b"abc") # revealed: bool
reveal_type(b"abc" is not b"ab") # revealed: Literal[True]
```

View File

@ -2660,6 +2660,51 @@ impl<'db> TypeInferenceBuilder<'db> {
self.infer_binary_type_comparison(left, op, KnownClass::Str.to_instance(self.db))
}
(Type::BytesLiteral(salsa_b1), Type::BytesLiteral(salsa_b2)) => {
let contains_subsequence = |needle: &[u8], haystack: &[u8]| {
if needle.is_empty() {
true
} else {
haystack
.windows(needle.len())
.any(|window| window == needle)
}
};
let b1 = salsa_b1.value(self.db).as_ref();
let b2 = salsa_b2.value(self.db).as_ref();
match op {
ast::CmpOp::Eq => Some(Type::BooleanLiteral(b1 == b2)),
ast::CmpOp::NotEq => Some(Type::BooleanLiteral(b1 != b2)),
ast::CmpOp::Lt => Some(Type::BooleanLiteral(b1 < b2)),
ast::CmpOp::LtE => Some(Type::BooleanLiteral(b1 <= b2)),
ast::CmpOp::Gt => Some(Type::BooleanLiteral(b1 > b2)),
ast::CmpOp::GtE => Some(Type::BooleanLiteral(b1 >= b2)),
ast::CmpOp::In => Some(Type::BooleanLiteral(contains_subsequence(b1, b2))),
ast::CmpOp::NotIn => Some(Type::BooleanLiteral(!contains_subsequence(b1, b2))),
ast::CmpOp::Is => {
if b1 == b2 {
Some(KnownClass::Bool.to_instance(self.db))
} else {
Some(Type::BooleanLiteral(false))
}
}
ast::CmpOp::IsNot => {
if b1 == b2 {
Some(KnownClass::Bool.to_instance(self.db))
} else {
Some(Type::BooleanLiteral(true))
}
}
}
}
(Type::BytesLiteral(_), _) => {
self.infer_binary_type_comparison(KnownClass::Bytes.to_instance(self.db), op, right)
}
(_, Type::BytesLiteral(_)) => {
self.infer_binary_type_comparison(left, op, KnownClass::Bytes.to_instance(self.db))
}
// Lookup the rich comparison `__dunder__` methods on instances
(Type::Instance(left_class_ty), Type::Instance(right_class_ty)) => match op {
ast::CmpOp::Lt => {