[red-knot] Fix inference for pow between two literal integers (#17161)

## Summary

Python `**` works differently to Rust `**`!

## Test Plan

Added an mdtest for various edge cases, and checked in the Python REPL
that we infer the correct type in all the new cases tested.
This commit is contained in:
Alex Waygood 2025-04-02 22:25:57 +01:00 committed by GitHub
parent 195bb433db
commit 28c68934a4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 22 additions and 9 deletions

View file

@ -55,6 +55,18 @@ def variable(x: int):
reveal_type(x**x) # revealed: @Todo(return type of overloaded function)
```
If the second argument is \<0, a `float` is returned at runtime. If the first argument is \<0 but
the second argument is >=0, an `int` is still returned:
```py
reveal_type(1**0) # revealed: Literal[1]
reveal_type(0**1) # revealed: Literal[0]
reveal_type(0**0) # revealed: Literal[1]
reveal_type((-1) ** 2) # revealed: Literal[1]
reveal_type(2 ** (-1)) # revealed: float
reveal_type((-1) ** (-1)) # revealed: float
```
## Division by Zero
This error is really outside the current Python type system, because e.g. `int.__truediv__` and

View file

@ -4615,16 +4615,17 @@ impl<'db> TypeInferenceBuilder<'db> {
.unwrap_or_else(|| KnownClass::Int.to_instance(self.db())),
),
(Type::IntLiteral(n), Type::IntLiteral(m), ast::Operator::Pow) => {
let m = u32::try_from(m);
Some(match m {
Ok(m) => n
.checked_pow(m)
(Type::IntLiteral(n), Type::IntLiteral(m), ast::Operator::Pow) => Some({
if m < 0 {
KnownClass::Float.to_instance(self.db())
} else {
u32::try_from(m)
.ok()
.and_then(|m| n.checked_pow(m))
.map(Type::IntLiteral)
.unwrap_or_else(|| KnownClass::Int.to_instance(self.db())),
Err(_) => KnownClass::Int.to_instance(self.db()),
})
.unwrap_or_else(|| KnownClass::Int.to_instance(self.db()))
}
}),
(Type::BytesLiteral(lhs), Type::BytesLiteral(rhs), ast::Operator::Add) => {
let bytes = [&**lhs.value(self.db()), &**rhs.value(self.db())].concat();