From 3629f3cada53663787095cd671d95ff85e27e9f4 Mon Sep 17 00:00:00 2001 From: WilliamK112 <164879897+WilliamK112@users.noreply.github.com> Date: Wed, 12 Aug 2026 16:11:23 +0800 Subject: [PATCH 1/2] test: reproduce invalid infer callable annotations Signed-off-by: WilliamK112 <164879897+WilliamK112@users.noreply.github.com> --- pyrefly/lib/commands/infer.rs | 47 +++++++++++++++++++++++++++++++++++ 1 file changed, 47 insertions(+) diff --git a/pyrefly/lib/commands/infer.rs b/pyrefly/lib/commands/infer.rs index 3901ad2e84..a8bd0e90e2 100644 --- a/pyrefly/lib/commands/infer.rs +++ b/pyrefly/lib/commands/infer.rs @@ -946,6 +946,53 @@ 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) +"#, + 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) +"#, + None, + ); + Ok(()) + } + #[test] fn test_imports() -> anyhow::Result<()> { let file_one = r#" From 12b7601074ef8320dfb947e11dfc3613f7f652c9 Mon Sep 17 00:00:00 2001 From: WilliamK112 <164879897+WilliamK112@users.noreply.github.com> Date: Wed, 12 Aug 2026 16:11:24 +0800 Subject: [PATCH 2/2] fix: render infer callables as Python annotations Signed-off-by: WilliamK112 <164879897+WilliamK112@users.noreply.github.com> --- pyrefly/lib/annotation.rs | 193 +++++++++++++++++++++++++++++++++ pyrefly/lib/commands/infer.rs | 45 +++++++- pyrefly/lib/lib.rs | 1 + pyrefly/lib/stubgen/extract.rs | 124 +-------------------- 4 files changed, 239 insertions(+), 124 deletions(-) create mode 100644 pyrefly/lib/annotation.rs diff --git a/pyrefly/lib/annotation.rs b/pyrefly/lib/annotation.rs new file mode 100644 index 0000000000..e169bcc9b2 --- /dev/null +++ b/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 { + 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::>(); + 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() + } +} diff --git a/pyrefly/lib/commands/infer.rs b/pyrefly/lib/commands/infer.rs index a8bd0e90e2..d2f8431ca8 100644 --- a/pyrefly/lib/commands/infer.rs +++ b/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,6 +292,8 @@ fn hint_to_string( stdlib: &Stdlib, enum_members: &dyn Fn(&Class) -> Option, 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(); @@ -281,10 +301,7 @@ fn hint_to_string( 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 { @@ -967,6 +984,14 @@ 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]): @@ -987,6 +1012,14 @@ 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, ); diff --git a/pyrefly/lib/lib.rs b/pyrefly/lib/lib.rs index 7d57335785..9bef9b1be2 100644 --- a/pyrefly/lib/lib.rs +++ b/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)] diff --git a/pyrefly/lib/stubgen/extract.rs b/pyrefly/lib/stubgen/extract.rs index e3675c06f6..478d8985b0 100644 --- a/pyrefly/lib/stubgen/extract.rs +++ b/pyrefly/lib/stubgen/extract.rs @@ -20,15 +20,9 @@ use pyrefly_python::ast::Ast; use pyrefly_python::module::Module; use pyrefly_python::short_identifier::ShortIdentifier; use pyrefly_python::sys_info::SysInfo; -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; use ruff_python_ast::BoolOp; use ruff_python_ast::Expr; @@ -44,6 +38,7 @@ use starlark_map::Hashed; use crate::alt::answers::Answers; use crate::alt::types::decorated_function::DecoratedFunction; +use crate::annotation::format_annotation; use crate::binding::binding::Binding; use crate::binding::binding::Key; use crate::binding::binding::KeyClassField; @@ -446,118 +441,11 @@ fn format_param_type(param: &Param, ctx: &mut ExtractionContext) -> Option Option { - if ty.is_any() { - ctx.uses_incomplete = true; - return Some("Incomplete".to_owned()); - } - // pyrefly's internal type display renders function and overload types using - // a non-Python `(args) -> ret` / `Overload[...]` syntax and never imports the - // referenced names. Translate them to valid `typing.Callable[...]` instead. - if let Some(s) = format_callable_type(ty, ctx) { - return Some(s); - } - if ty.any(|sub_type| matches!(sub_type, Type::SelfType(_))) { - ctx.typing_imports.insert("Self"); - } - if ty.any(|sub_type| matches!(sub_type, Type::Literal(_))) { - ctx.typing_imports.insert("Literal"); - } - let mut display = TypeDisplayContext::new(&[ty]); - display.render_self_type_as_self(); - display.strip_library_schemas(); - let s = display.display(ty).to_string(); - if s.contains("@") || s.contains("Unknown") { - ctx.uses_incomplete = true; - return Some("Incomplete".to_owned()); - } - Some(s) -} - -/// Render a callable-typed value as a valid `typing.Callable[...]` annotation, -/// or `None` if `ty` is not a callable type. -fn format_callable_type(ty: &Type, ctx: &mut ExtractionContext) -> Option { - match ty { - Type::Function(func) => Some(callable_from_signature(&func.signature, ctx)), - Type::Callable(callable) => Some(callable_from_signature(callable, ctx)), - Type::BoundMethod(method) => match &method.func { - // Drop the bound `self`/`cls` parameter before rendering. - BoundMethodType::Function(func) => { - let sig = func.signature.strip_first_param().expect( - "BoundMethod::Function should always have at least a self/cls parameter", - ); - Some(callable_from_signature(&sig, ctx)) - } - // Generic methods: parameters can't be faithfully expressed. - BoundMethodType::Forall(forall) => { - Some(callable_ellipsis(&forall.body.signature.ret, ctx)) - } - BoundMethodType::Overload(overload) => Some(format_overload(overload, ctx)), - }, - Type::Overload(overload) => Some(format_overload(overload, ctx)), - // A generic (`Forall`) function/callable: parameters carry type - // variables we can't faithfully express, so elide them. - Type::Forall(forall) => match &forall.body { - Forallable::Function(func) => Some(callable_ellipsis(&func.signature.ret, ctx)), - Forallable::Callable(callable) => Some(callable_ellipsis(&callable.ret, ctx)), - Forallable::TypeAlias(_) => None, - }, - _ => None, - } -} - -/// Render `Callable[[A, B], Ret]` when the parameter list is a plain sequence of -/// required positional parameters, otherwise `Callable[..., Ret]` (e.g. when it -/// contains `*args`, `**kwargs`, keyword-only, or optional parameters). The -/// explicit-argument form requires exactly those arguments, so an optional param -/// would over-constrain callers; eliding to `...` avoids that. -fn callable_from_signature(sig: &Callable, ctx: &mut ExtractionContext) -> String { - ctx.typing_imports.insert("Callable"); - let ret = format_type(&sig.ret, ctx).expect("format_type always returns Some"); - match &sig.params { - Params::List(params) - if params.items().iter().all(|p| { - matches!( - p, - Param::PosOnly(_, _, Required::Required) | Param::Pos(_, _, Required::Required) - ) - }) => - { - let rendered: Vec = params - .items() - .iter() - .map(|p| format_type(p.as_type(), ctx).expect("format_type always returns Some")) - .collect(); - format!("Callable[[{}], {}]", rendered.join(", "), ret) - } - _ => format!("Callable[..., {ret}]"), - } -} - -/// `Callable[..., Ret]` for cases where the parameters can't be expressed. -fn callable_ellipsis(ret: &Type, ctx: &mut ExtractionContext) -> String { - ctx.typing_imports.insert("Callable"); - let ret = format_type(ret, ctx).expect("format_type always returns Some"); - format!("Callable[..., {ret}]") -} - -/// Render an overload set. If every signature shares the same return type we -/// can render `Callable[..., Ret]`; otherwise there is no single faithful -/// return type, so fall back to `Callable[..., Incomplete]` (still a callable). -fn format_overload(overload: &Overload, ctx: &mut ExtractionContext) -> String { - let ret = |sig: &OverloadType| -> Type { - match sig { - OverloadType::Function(f) => f.signature.ret.clone(), - OverloadType::Forall(forall) => forall.body.signature.ret.clone(), - } - }; - let first = ret(overload.signatures.first()); - if overload.signatures.iter().all(|s| ret(s) == first) { - callable_ellipsis(&first, ctx) - } else { - ctx.typing_imports.insert("Callable"); - ctx.uses_incomplete = true; - "Callable[..., Incomplete]".to_owned() - } + Some(format_annotation( + ty, + &mut ctx.typing_imports, + &mut ctx.uses_incomplete, + )) } /// Uses source text for simple literals, `...` for everything else.