[red-knot] Add CallableTypeFromFunction special form

crates/red_knot_python_semantic/resources/mdtest/type_api.md

@@ -393,3 +393,42 @@ def type_of_annotation() -> None:
 # error: "Special form `knot_extensions.TypeOf` expected exactly one type parameter"
 t: TypeOf[int, str, bytes]
 ```
+
+## `CallableTypeFromFunction`
+
+The `CallableTypeFromFunction` special form can be used to extract the type of a function literal as
+a callable type. This can be used to get the externally-visibly signature of the function, which can
+then be used to test various type properties.
+
+It accepts a single type parameter which is expected to be a function literal.
+
+```py
+from knot_extensions import CallableTypeFromFunction
+
+def f1():
+    return
+
+def f2() -> int:
+    return 1
+
+def f3(x: int, y: str) -> None:
+    return
+
+# error: [invalid-type-form] "Special form `knot_extensions.CallableTypeFromFunction` expected exactly one type parameter"
+c1: CallableTypeFromFunction[f1, f2]
+# error: [invalid-type-form] "Expected the first argument to `knot_extensions.CallableTypeFromFunction` to be a function literal, but got `Literal[int]`"
+c2: CallableTypeFromFunction[int]
+```
+
+Using it in annotation to reveal the signature of the function:
+
+```py
+def _(
+    c1: CallableTypeFromFunction[f1],
+    c2: CallableTypeFromFunction[f2],
+    c3: CallableTypeFromFunction[f3],
+) -> None:
+    reveal_type(c1)  # revealed: () -> Unknown
+    reveal_type(c2)  # revealed: () -> int
+    reveal_type(c3)  # revealed: (x: int, y: str) -> None
+```

crates/red_knot_python_semantic/src/types.rs

@@ -4300,6 +4300,17 @@ impl<'db> FunctionType<'db> {
         })
     }
 
+    /// Convert the `FunctionType` into a [`Type::Callable`].
+    ///
+    /// Returns `None` if the function is overloaded. This powers the `CallableTypeFromFunction`
+    /// special form from the `knot_extensions` module.
+    pub(crate) fn into_callable_type(self, db: &'db dyn Db) -> Option<Type<'db>> {
+        // TODO: Add support for overloaded callables; return `Type`, not `Option<Type>`.
+        Some(Type::Callable(CallableType::General(
+            GeneralCallableType::new(db, self.signature(db).as_single()?.clone()),
+        )))
+    }
+
     /// Typed externally-visible signature for this function.
     ///
     /// This is the signature as seen by external callers, possibly modified by decorators and/or

crates/red_knot_python_semantic/src/types/class.rs

@@ -1434,6 +1434,8 @@ pub enum KnownInstanceType<'db> {
     Intersection,
     /// The symbol `knot_extensions.TypeOf`
     TypeOf,
+    /// The symbol `knot_extensions.CallableTypeFromFunction`
+    CallableTypeFromFunction,
 
     // Various special forms, special aliases and type qualifiers that we don't yet understand
     // (all currently inferred as TODO in most contexts):
@@ -1495,7 +1497,8 @@ impl<'db> KnownInstanceType<'db> {
             | Self::AlwaysFalsy
             | Self::Not
             | Self::Intersection
-            | Self::TypeOf => Truthiness::AlwaysTrue,
+            | Self::TypeOf
+            | Self::CallableTypeFromFunction => Truthiness::AlwaysTrue,
         }
     }
 
@@ -1542,6 +1545,7 @@ impl<'db> KnownInstanceType<'db> {
             Self::Not => "knot_extensions.Not",
             Self::Intersection => "knot_extensions.Intersection",
             Self::TypeOf => "knot_extensions.TypeOf",
+            Self::CallableTypeFromFunction => "knot_extensions.CallableTypeFromFunction",
         }
     }
 
@@ -1585,6 +1589,7 @@ impl<'db> KnownInstanceType<'db> {
             Self::TypeOf => KnownClass::SpecialForm,
             Self::Not => KnownClass::SpecialForm,
             Self::Intersection => KnownClass::SpecialForm,
+            Self::CallableTypeFromFunction => KnownClass::SpecialForm,
             Self::Unknown => KnownClass::Object,
             Self::AlwaysTruthy => KnownClass::Object,
             Self::AlwaysFalsy => KnownClass::Object,
@@ -1649,6 +1654,7 @@ impl<'db> KnownInstanceType<'db> {
             "Not" => Self::Not,
             "Intersection" => Self::Intersection,
             "TypeOf" => Self::TypeOf,
+            "CallableTypeFromFunction" => Self::CallableTypeFromFunction,
             _ => return None,
         };
 
@@ -1704,7 +1710,8 @@ impl<'db> KnownInstanceType<'db> {
             | Self::AlwaysFalsy
             | Self::Not
             | Self::Intersection
-            | Self::TypeOf => module.is_knot_extensions(),
+            | Self::TypeOf
+            | Self::CallableTypeFromFunction => module.is_knot_extensions(),
         }
     }
 }

crates/red_knot_python_semantic/src/types/class_base.rs

@@ -103,6 +103,7 @@ impl<'db> ClassBase<'db> {
                 | KnownInstanceType::Not
                 | KnownInstanceType::Intersection
                 | KnownInstanceType::TypeOf
+                | KnownInstanceType::CallableTypeFromFunction
                 | KnownInstanceType::AlwaysTruthy
                 | KnownInstanceType::AlwaysFalsy => None,
                 KnownInstanceType::Unknown => Some(Self::unknown()),

crates/red_knot_python_semantic/src/types/infer.rs

@@ -6281,6 +6281,44 @@ impl<'db> TypeInferenceBuilder<'db> {
                     argument_type
                 }
             },
+            KnownInstanceType::CallableTypeFromFunction => match arguments_slice {
+                ast::Expr::Tuple(_) => {
+                    self.context.report_lint(
+                        &INVALID_TYPE_FORM,
+                        subscript,
+                        format_args!(
+                            "Special form `{}` expected exactly one type parameter",
+                            known_instance.repr(self.db())
+                        ),
+                    );
+                    Type::unknown()
+                }
+                _ => {
+                    let argument_type = self.infer_expression(arguments_slice);
+                    let Some(function_type) = argument_type.into_function_literal() else {
+                        self.context.report_lint(
+                            &INVALID_TYPE_FORM,
+                            arguments_slice,
+                            format_args!(
+                                "Expected the first argument to `{}` to be a function literal, but got `{}`",
+                                known_instance.repr(self.db()),
+                                argument_type.display(self.db())
+                            ),
+                        );
+                        return Type::unknown();
+                    };
+                    function_type
+                        .into_callable_type(self.db())
+                        .unwrap_or_else(|| {
+                            self.context.report_lint(
+                                &INVALID_TYPE_FORM,
+                                arguments_slice,
+                                format_args!("Overloaded function literal is not yet supported"),
+                            );
+                            Type::unknown()
+                        })
+                }
+            },
 
             // TODO: Generics
             KnownInstanceType::ChainMap => {

crates/red_knot_python_semantic/src/types/signatures.rs

@@ -41,6 +41,14 @@ impl<'db> CallableSignature<'db> {
         CallableSignature::Overloaded(overloads.into())
     }
 
+    /// Returns the [`Signature`] if this is a non-overloaded callable, [None] otherwise.
+    pub(crate) fn as_single(&self) -> Option<&Signature<'db>> {
+        match self {
+            CallableSignature::Single(signature) => Some(signature),
+            CallableSignature::Overloaded(_) => None,
+        }
+    }
+
     pub(crate) fn iter(&self) -> std::slice::Iter<Signature<'db>> {
         match self {
             CallableSignature::Single(signature) => std::slice::from_ref(signature).iter(),

crates/red_knot_python_semantic/src/types/type_ordering.rs

@@ -220,6 +220,9 @@ pub(super) fn union_elements_ordering<'db>(left: &Type<'db>, right: &Type<'db>)
                 (KnownInstanceType::TypeOf, _) => Ordering::Less,
                 (_, KnownInstanceType::TypeOf) => Ordering::Greater,
 
+                (KnownInstanceType::CallableTypeFromFunction, _) => Ordering::Less,
+                (_, KnownInstanceType::CallableTypeFromFunction) => Ordering::Greater,
+
                 (KnownInstanceType::Unpack, _) => Ordering::Less,
                 (_, KnownInstanceType::Unpack) => Ordering::Greater,
 

crates/red_knot_vendored/knot_extensions/knot_extensions.pyi

@@ -12,6 +12,7 @@ AlwaysFalsy = object()
 Not: _SpecialForm
 Intersection: _SpecialForm
 TypeOf: _SpecialForm
+CallableTypeFromFunction: _SpecialForm
 
 # Predicates on types
 #