From 28c68934a42ddb6dec4ad5324a0546d13a35771d Mon Sep 17 00:00:00 2001
From: Alex Waygood <Alex.Waygood@Gmail.com>
Date: Wed, 2 Apr 2025 22:25:57 +0100
Subject: [PATCH] [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.
---
 .../resources/mdtest/binary/integers.md       | 12 ++++++++++++
 .../src/types/infer.rs                        | 19 ++++++++++---------
 2 files changed, 22 insertions(+), 9 deletions(-)

diff --git a/crates/red_knot_python_semantic/resources/mdtest/binary/integers.md b/crates/red_knot_python_semantic/resources/mdtest/binary/integers.md
index a041bbf47a..a4172c821d 100644
--- a/crates/red_knot_python_semantic/resources/mdtest/binary/integers.md
+++ b/crates/red_knot_python_semantic/resources/mdtest/binary/integers.md
@@ -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
diff --git a/crates/red_knot_python_semantic/src/types/infer.rs b/crates/red_knot_python_semantic/src/types/infer.rs
index 243f68a83f..b0766921fe 100644
--- a/crates/red_knot_python_semantic/src/types/infer.rs
+++ b/crates/red_knot_python_semantic/src/types/infer.rs
@@ -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();