From 58d5fe982ed8fd248f04c16955eb59c9b49b540b Mon Sep 17 00:00:00 2001
From: Dhruv Manilawala <dhruvmanila@gmail.com>
Date: Thu, 13 Mar 2025 08:16:51 +0530
Subject: [PATCH] [red-knot] Check gradual equivalence between callable types
 (#16634)

---
 .../is_gradual_equivalent_to.md               | 63 +++++++++++++++++++
 crates/red_knot_python_semantic/src/types.rs  | 47 ++++++++++++++
 2 files changed, 110 insertions(+)

diff --git a/crates/red_knot_python_semantic/resources/mdtest/type_properties/is_gradual_equivalent_to.md b/crates/red_knot_python_semantic/resources/mdtest/type_properties/is_gradual_equivalent_to.md
index 5dab6ebeab..795e7a5822 100644
--- a/crates/red_knot_python_semantic/resources/mdtest/type_properties/is_gradual_equivalent_to.md
+++ b/crates/red_knot_python_semantic/resources/mdtest/type_properties/is_gradual_equivalent_to.md
@@ -62,4 +62,67 @@ static_assert(not is_gradual_equivalent_to(tuple[str, int], tuple[str, int, byte
 static_assert(not is_gradual_equivalent_to(tuple[str, int], tuple[int, str]))
 ```
 
+## Callable
+
+```py
+from knot_extensions import Unknown, CallableTypeFromFunction, is_gradual_equivalent_to, static_assert
+from typing import Any, Callable
+
+static_assert(is_gradual_equivalent_to(Callable[..., int], Callable[..., int]))
+static_assert(is_gradual_equivalent_to(Callable[..., Any], Callable[..., Unknown]))
+static_assert(is_gradual_equivalent_to(Callable[[int, Any], None], Callable[[int, Unknown], None]))
+
+static_assert(not is_gradual_equivalent_to(Callable[[int, Any], None], Callable[[Any, int], None]))
+static_assert(not is_gradual_equivalent_to(Callable[[int, str], None], Callable[[int, str, bytes], None]))
+static_assert(not is_gradual_equivalent_to(Callable[..., None], Callable[[], None]))
+```
+
+A function with no explicit return type should be gradual equivalent to a callable with a return
+type of `Any`.
+
+```py
+def f1():
+    return
+
+static_assert(is_gradual_equivalent_to(CallableTypeFromFunction[f1], Callable[[], Any]))
+```
+
+And, similarly for parameters with no annotations.
+
+```py
+def f2(a, b) -> None:
+    return
+
+static_assert(is_gradual_equivalent_to(CallableTypeFromFunction[f2], Callable[[Any, Any], None]))
+```
+
+Additionally, as per the spec, a function definition that includes both `*args` and `**kwargs`
+parameter that are annotated as `Any` or kept unannotated should be gradual equivalent to a callable
+with `...` as the parameter type.
+
+```py
+def variadic_without_annotation(*args, **kwargs):
+    return
+
+def variadic_with_annotation(*args: Any, **kwargs: Any) -> Any:
+    return
+
+static_assert(is_gradual_equivalent_to(CallableTypeFromFunction[variadic_without_annotation], Callable[..., Any]))
+static_assert(is_gradual_equivalent_to(CallableTypeFromFunction[variadic_with_annotation], Callable[..., Any]))
+```
+
+But, a function with either `*args` or `**kwargs` is not gradual equivalent to a callable with `...`
+as the parameter type.
+
+```py
+def variadic_args(*args):
+    return
+
+def variadic_kwargs(**kwargs):
+    return
+
+static_assert(not is_gradual_equivalent_to(CallableTypeFromFunction[variadic_args], Callable[..., Any]))
+static_assert(not is_gradual_equivalent_to(CallableTypeFromFunction[variadic_kwargs], Callable[..., Any]))
+```
+
 [materializations]: https://typing.readthedocs.io/en/latest/spec/glossary.html#term-materialize
diff --git a/crates/red_knot_python_semantic/src/types.rs b/crates/red_knot_python_semantic/src/types.rs
index 84035288f8..2d5f05740a 100644
--- a/crates/red_knot_python_semantic/src/types.rs
+++ b/crates/red_knot_python_semantic/src/types.rs
@@ -956,6 +956,11 @@ impl<'db> Type<'db> {
                 first.is_gradual_equivalent_to(db, second)
             }
 
+            (
+                Type::Callable(CallableType::General(first)),
+                Type::Callable(CallableType::General(second)),
+            ) => first.is_gradual_equivalent_to(db, second),
+
             _ => false,
         }
     }
@@ -4575,6 +4580,48 @@ impl<'db> GeneralCallableType<'db> {
             .return_ty
             .is_some_and(|return_type| return_type.is_fully_static(db))
     }
+
+    /// Return `true` if `self` has exactly the same set of possible static materializations as
+    /// `other` (if `self` represents the same set of possible sets of possible runtime objects as
+    /// `other`).
+    pub(crate) fn is_gradual_equivalent_to(self, db: &'db dyn Db, other: Self) -> bool {
+        let self_signature = self.signature(db);
+        let other_signature = other.signature(db);
+
+        if self_signature.parameters().len() != other_signature.parameters().len() {
+            return false;
+        }
+
+        // Check gradual equivalence between the two optional types. In the context of a callable
+        // type, the `None` type represents an `Unknown` type.
+        let are_optional_types_gradually_equivalent =
+            |self_type: Option<Type<'db>>, other_type: Option<Type<'db>>| {
+                self_type
+                    .unwrap_or(Type::unknown())
+                    .is_gradual_equivalent_to(db, other_type.unwrap_or(Type::unknown()))
+            };
+
+        if !are_optional_types_gradually_equivalent(
+            self_signature.return_ty,
+            other_signature.return_ty,
+        ) {
+            return false;
+        }
+
+        // N.B. We don't need to explicitly check for the use of gradual form (`...`) in the
+        // parameters because it is internally represented by adding `*Any` and `**Any` to the
+        // parameter list.
+        self_signature
+            .parameters()
+            .iter()
+            .zip(other_signature.parameters().iter())
+            .all(|(self_param, other_param)| {
+                are_optional_types_gradually_equivalent(
+                    self_param.annotated_type(),
+                    other_param.annotated_type(),
+                )
+            })
+    }
 }
 
 /// A type that represents callable objects.