Fix infer callable annotations

pyrefly/lib/annotation.rs

@@ -0,0 +1,193 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+use std::collections::BTreeSet;
+
+use pyrefly_types::callable::Callable;
+use pyrefly_types::callable::Param;
+use pyrefly_types::callable::Params;
+use pyrefly_types::callable::Required;
+use pyrefly_types::display::TypeDisplayContext;
+use pyrefly_types::types::BoundMethodType;
+use pyrefly_types::types::Forallable;
+use pyrefly_types::types::Overload;
+use pyrefly_types::types::OverloadType;
+use pyrefly_types::types::Type;
+
+/// Render an internal type as valid Python annotation syntax.
+pub(crate) fn format_annotation(
+    ty: &Type,
+    typing_imports: &mut BTreeSet<&'static str>,
+    uses_incomplete: &mut bool,
+) -> String {
+    if ty.is_any() {
+        *uses_incomplete = true;
+        return "Incomplete".to_owned();
+    }
+    // Internal function and overload displays use `(args) -> ret` / `Overload[...]`,
+    // which are not valid Python annotations.
+    if let Some(annotation) = format_callable_type(ty, typing_imports, uses_incomplete) {
+        return annotation;
+    }
+    // Nested internal callable displays would also produce invalid Python, but
+    // rewriting them would require reconstructing the enclosing type. Callers
+    // can either emit `Incomplete` or decline to add the annotation.
+    if ty.any(is_callable_type) {
+        *uses_incomplete = true;
+        return "Incomplete".to_owned();
+    }
+    if ty.any(|sub_type| matches!(sub_type, Type::SelfType(_))) {
+        typing_imports.insert("Self");
+    }
+    if ty.any(|sub_type| matches!(sub_type, Type::Literal(_))) {
+        typing_imports.insert("Literal");
+    }
+    let mut display = TypeDisplayContext::new(&[ty]);
+    display.render_self_type_as_self();
+    display.strip_library_schemas();
+    let annotation = display.display(ty).to_string();
+    if annotation.contains('@') || annotation.contains("Unknown") {
+        *uses_incomplete = true;
+        "Incomplete".to_owned()
+    } else {
+        annotation
+    }
+}
+
+/// Whether the internal display for this type uses callable-only syntax.
+fn is_callable_type(ty: &Type) -> bool {
+    match ty {
+        Type::Function(_) | Type::Callable(_) | Type::BoundMethod(_) | Type::Overload(_) => true,
+        Type::Forall(forall) => matches!(
+            &forall.body,
+            Forallable::Function(_) | Forallable::Callable(_)
+        ),
+        _ => false,
+    }
+}
+
+/// Render a callable-typed value as `typing.Callable[...]`.
+fn format_callable_type(
+    ty: &Type,
+    typing_imports: &mut BTreeSet<&'static str>,
+    uses_incomplete: &mut bool,
+) -> Option<String> {
+    match ty {
+        Type::Function(func) => Some(callable_from_signature(
+            &func.signature,
+            typing_imports,
+            uses_incomplete,
+        )),
+        Type::Callable(callable) => Some(callable_from_signature(
+            callable,
+            typing_imports,
+            uses_incomplete,
+        )),
+        Type::BoundMethod(method) => match &method.func {
+            BoundMethodType::Function(func) => {
+                let signature = func.signature.strip_first_param().expect(
+                    "BoundMethod::Function should always have at least a self/cls parameter",
+                );
+                Some(callable_from_signature(
+                    &signature,
+                    typing_imports,
+                    uses_incomplete,
+                ))
+            }
+            BoundMethodType::Forall(forall) => Some(callable_ellipsis(
+                &forall.body.signature.ret,
+                typing_imports,
+                uses_incomplete,
+            )),
+            BoundMethodType::Overload(overload) => {
+                Some(format_overload(overload, typing_imports, uses_incomplete))
+            }
+        },
+        Type::Overload(overload) => {
+            Some(format_overload(overload, typing_imports, uses_incomplete))
+        }
+        Type::Forall(forall) => match &forall.body {
+            Forallable::Function(func) => Some(callable_ellipsis(
+                &func.signature.ret,
+                typing_imports,
+                uses_incomplete,
+            )),
+            Forallable::Callable(callable) => Some(callable_ellipsis(
+                &callable.ret,
+                typing_imports,
+                uses_incomplete,
+            )),
+            Forallable::TypeAlias(_) => None,
+        },
+        _ => None,
+    }
+}
+
+/// Preserve required positional parameters and elide signatures that `Callable` cannot express.
+fn callable_from_signature(
+    signature: &Callable,
+    typing_imports: &mut BTreeSet<&'static str>,
+    uses_incomplete: &mut bool,
+) -> String {
+    typing_imports.insert("Callable");
+    let ret = format_annotation(&signature.ret, typing_imports, uses_incomplete);
+    match &signature.params {
+        Params::List(params)
+            if params.items().iter().all(|param| {
+                matches!(
+                    param,
+                    Param::PosOnly(_, _, Required::Required) | Param::Pos(_, _, Required::Required)
+                )
+            }) =>
+        {
+            let rendered = params
+                .items()
+                .iter()
+                .map(|param| format_annotation(param.as_type(), typing_imports, uses_incomplete))
+                .collect::<Vec<_>>();
+            format!("Callable[[{}], {}]", rendered.join(", "), ret)
+        }
+        _ => format!("Callable[..., {ret}]"),
+    }
+}
+
+/// Render a callable whose parameter list cannot be faithfully expressed.
+fn callable_ellipsis(
+    ret: &Type,
+    typing_imports: &mut BTreeSet<&'static str>,
+    uses_incomplete: &mut bool,
+) -> String {
+    typing_imports.insert("Callable");
+    let ret = format_annotation(ret, typing_imports, uses_incomplete);
+    format!("Callable[..., {ret}]")
+}
+
+/// Render overloads with a shared return type, or an incomplete callable otherwise.
+fn format_overload(
+    overload: &Overload,
+    typing_imports: &mut BTreeSet<&'static str>,
+    uses_incomplete: &mut bool,
+) -> String {
+    let ret = |signature: &OverloadType| -> Type {
+        match signature {
+            OverloadType::Function(function) => function.signature.ret.clone(),
+            OverloadType::Forall(forall) => forall.body.signature.ret.clone(),
+        }
+    };
+    let first = ret(overload.signatures.first());
+    if overload
+        .signatures
+        .iter()
+        .all(|signature| ret(signature) == first)
+    {
+        callable_ellipsis(&first, typing_imports, uses_incomplete)
+    } else {
+        typing_imports.insert("Callable");
+        *uses_incomplete = true;
+        "Callable[..., Incomplete]".to_owned()
+    }
+}

pyrefly/lib/commands/infer.rs

@@ -5,6 +5,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+use std::collections::BTreeSet;
 use std::collections::HashSet;
 use std::path::Path;
 
@@ -23,6 +24,7 @@ use ruff_python_ast::helpers::is_docstring_stmt;
 use ruff_text_size::Ranged;
 use ruff_text_size::TextSize;
 
+use crate::annotation::format_annotation;
 use crate::commands::check::Handles;
 use crate::commands::config_finder::ConfigConfigurerWrapper;
 use crate::commands::files::FilesArgs;
@@ -32,7 +34,6 @@ use crate::state::lsp::AnnotationKind;
 use crate::state::require::Require;
 use crate::state::state::State;
 use crate::types::class::Class;
-use crate::types::display::TypeDisplayContext;
 use crate::types::heap::TypeHeap;
 use crate::types::simplify::unions_with_literals;
 use crate::types::stdlib::Stdlib;
@@ -220,7 +221,16 @@ fn format_hints(
         if contains_self_type {
             hint_imports.push((ModuleName::typing(), "Self".to_owned()));
         }
-        let formatted_hint = hint_to_string(hint, stdlib, enum_members, heap);
+        let mut typing_imports = BTreeSet::new();
+        let mut uses_incomplete = false;
+        let formatted_hint = hint_to_string(
+            hint,
+            stdlib,
+            enum_members,
+            heap,
+            &mut typing_imports,
+            &mut uses_incomplete,
+        );
         // TODO: Put these behind a flag
         if formatted_hint.contains("Any") {
             continue;
@@ -237,12 +247,20 @@ fn format_hints(
         if formatted_hint.contains("Overload") {
             continue;
         }
+        if uses_incomplete {
+            continue;
+        }
         if formatted_hint == "None" && kind == AnnotationKind::Parameter {
             continue;
         }
         if !is_container && kind == AnnotationKind::Variable {
             continue;
         }
+        needed_imports.extend(
+            typing_imports
+                .into_iter()
+                .map(|name| (ModuleName::typing(), name.to_owned())),
+        );
         // Only record imports for types that pass all filters above
         needed_imports.extend(hint_imports);
         match kind {
@@ -274,17 +292,16 @@ fn hint_to_string(
     stdlib: &Stdlib,
     enum_members: &dyn Fn(&Class) -> Option<usize>,
     heap: &TypeHeap,
+    typing_imports: &mut BTreeSet<&'static str>,
+    uses_incomplete: &mut bool,
 ) -> String {
     let hint = hint.promote_implicit_literals(stdlib);
     let hint = hint.explicit_any().clean_var();
     let hint = match hint {
         Type::Union(u) => unions_with_literals(u.members, stdlib, enum_members, heap),
         _ => hint,
     };
-    let mut ctx = TypeDisplayContext::new(&[&hint]);
-    ctx.render_self_type_as_self();
-    ctx.strip_library_schemas();
-    ctx.display(&hint).to_string()
+    format_annotation(&hint, typing_imports, uses_incomplete)
 }
 
 impl InferArgs {
@@ -946,6 +963,69 @@ class C:
         Ok(())
     }
 
+    #[test]
+    fn test_callable_annotations_use_python_syntax() -> anyhow::Result<()> {
+        assert_annotations(
+            r#"def call_it(fn):
+    return fn()
+
+call_it(lambda: 0)
+
+def make_formatter():
+    def format_one(n: int) -> str:
+        return str(n)
+    return format_one
+
+class Runner:
+    def run(self) -> None:
+        pass
+
+def swallow(fn):
+    fn()
+
+swallow(Runner().run)
+
+def maybe_callback(flag: bool):
+    if flag:
+        return lambda: 1
+    return None
+
+def callbacks():
+    return [lambda: 1]
+"#,
+            r#"from typing import Callable
+def call_it(fn: Callable[[], int]):
+    return fn()
+
+call_it(lambda: 0)
+
+def make_formatter() -> Callable[[int], str]:
+    def format_one(n: int) -> str:
+        return str(n)
+    return format_one
+
+class Runner:
+    def run(self) -> None:
+        pass
+
+def swallow(fn: Callable[[], None]) -> None:
+    fn()
+
+swallow(Runner().run)
+
+def maybe_callback(flag: bool):
+    if flag:
+        return lambda: 1
+    return None
+
+def callbacks():
+    return [lambda: 1]
+"#,
+            None,
+        );
+        Ok(())
+    }
+
     #[test]
     fn test_imports() -> anyhow::Result<()> {
         let file_one = r#"

pyrefly/lib/lib.rs

@@ -25,6 +25,7 @@
 #![deny(clippy::trivially_copy_pass_by_ref)]
 
 pub mod alt;
+mod annotation;
 pub mod binding;
 #[cfg(not(target_arch = "wasm32"))]
 #[doc(hidden)]