map (...) -> R to Callable[..., R] and strip stray stub bodies

Author: KotlinIsland

crates/by_transforms/src/lib.rs

@@ -487,12 +487,16 @@ pub fn reverse_transpile(source: &str, config: &Config) -> Result<String, String
     let mut super_kw_rev = reverse_transforms::super_keyword::SuperKeywordReverse::new(src);
     let mut anon_named_tuple_rev =
         reverse_transforms::anon_named_tuple::AnonNamedTupleReverse::new(src, module.suite());
-    let mut empty_decls = reverse_transforms::empty_declarations::EmptyDeclarations::new();
+    let mut empty_decls =
+        reverse_transforms::empty_declarations::EmptyDeclarations::new(config.is_stub);
     let mut literal_types = reverse_transforms::literal_types::LiteralReverse::new(src, &model);
     let mut subscript = reverse_transforms::subscript::SubscriptReverse::new(src, &model);
     let mut indent_string = reverse_transforms::dedent_string::IndentString::new(src);
     let mut constraints = reverse_transforms::constraints::ConstraintsReverse::new();
-    let mut callable = reverse_transforms::callable::CallableReverse::new(src, &model);
+    let mut callable = {
+        let c = reverse_transforms::callable::CallableReverse::new(src, &model);
+        if config.is_stub { c.stub() } else { c }
+    };
     let mut intersection = reverse_transforms::intersection::IntersectionReverse::new(src, &model);
     let mut not_rev = reverse_transforms::not_type::NotTypeReverse::new(src, &model);
     let mut type_is_rev = reverse_transforms::type_is::TypeIsReverse::new(src, &model);
@@ -525,21 +529,24 @@ pub fn reverse_transpile(source: &str, config: &Config) -> Result<String, String
         auto_quote_rev.visit_stmt(stmt);
         compat_rev.visit_stmt(stmt);
         none_chain_rev.visit_stmt(stmt);
+        // `callable` rewrites callable annotations to the arrow form. it runs
+        // for stubs too, but in a restricted "stub" mode (set above) that only
+        // touches the gradual `Callable[..., R]` form — the `Callable[[A, B],
+        // R]` list form is left intact, since ty's native basedpython parser
+        // can't carry `Unpack[Ts]`/`*Ts` through the arrow and stubs would
+        // lose generic callable info
+        callable.visit_stmt(stmt);
         // skip transforms that change runtime/display semantics when
         // rewriting stubs:
         //  - `literal_types` strips `Literal[...]` to bare literals, but a
         //    bare `1 | 2` in a `TypeAlias = ...` RHS evaluates at runtime
         //    as integer OR (= `3`) rather than `Literal[1, 2]`
-        //  - `callable` rewrites `Callable[[A, B], R]` to `(A, B) -> R`;
-        //    ty's native basedpython parser doesn't handle `Unpack[Ts]` and
-        //    `*Ts` inside the arrow form, so stubs lose generic callable info
         //  - `typing_redirect` rewrites `typing_extensions` imports, but
         //    stubs use them deliberately for version-aware re-exports
         //  - `generics` turns `X: TypeAlias = T` into PEP 695 `type X = T`,
         //    which resolves lazily and changes alias display in diagnostics
         if !config.is_stub {
             literal_types.visit_stmt(stmt);
-            callable.visit_stmt(stmt);
             typing_redirect_rev.visit_stmt(stmt);
             generics_rev.visit_stmt(stmt);
         }

crates/by_transforms/src/reverse_transforms/callable.rs

@@ -7,14 +7,20 @@
 
 use ruff_diagnostics::{Edit, Fix};
 use ruff_python_ast::visitor::Visitor;
-use ruff_python_ast::{Expr, Stmt};
+use ruff_python_ast::{Expr, ExprSubscript, Stmt};
 use ruff_text_size::{Ranged, TextRange};
 
 use crate::type_info::TypeInfo;
 
 pub(crate) struct CallableReverse<'src> {
     source: &'src str,
     types: &'src dyn TypeInfo,
+    /// in stub mode, the `Callable[[A, B], R]` list form is left intact —
+    /// ty's basedpython parser can't carry `Unpack[Ts]`/`*Ts` through the
+    /// arrow form, so stubs would lose generic callable info. the gradual
+    /// `Callable[..., R]` form has no parameter list to lose and is always
+    /// rewritten to `(...) -> R`
+    stub: bool,
     pub(crate) edits: Vec<Fix>,
 }
 
@@ -23,10 +29,16 @@ impl<'src> CallableReverse<'src> {
         Self {
             source,
             types,
+            stub: false,
             edits: Vec::new(),
         }
     }
 
+    pub(crate) fn stub(mut self) -> Self {
+        self.stub = true;
+        self
+    }
+
     fn src(&self, range: TextRange) -> &str {
         &self.source[usize::from(range.start())..usize::from(range.end())]
     }
@@ -51,9 +63,27 @@ impl<'src> CallableReverse<'src> {
                 if t.parenthesized || t.elts.len() != 2 {
                     return None;
                 }
+                let ret = &t.elts[1];
+                // `Callable[..., R]` — "any arguments" — reverses to `(...) -> R`.
+                // safe in stub mode: there is no parameter list to lose
+                if matches!(&t.elts[0], Expr::EllipsisLiteral(_)) {
+                    let ret_str = self
+                        .rewrite(ret)
+                        .unwrap_or_else(|| self.src(ret.range()).to_owned());
+                    return Some(format!("(...) -> {ret_str}"));
+                }
+                // list form: leave the `Callable[...]` wrapper intact in stub
+                // mode but still recurse so any nested `Callable[..., R]` is
+                // converted
+                if self.stub {
+                    return self.rewrite_subscript_children(s);
+                }
                 let Expr::List(args_list) = &t.elts[0] else {
                     return None;
                 };
+                let ret_str = self
+                    .rewrite(ret)
+                    .unwrap_or_else(|| self.src(ret.range()).to_owned());
                 let args_str = args_list
                     .elts
                     .iter()
@@ -63,10 +93,6 @@ impl<'src> CallableReverse<'src> {
                     })
                     .collect::<Vec<_>>()
                     .join(", ");
-                let ret = &t.elts[1];
-                let ret_str = self
-                    .rewrite(ret)
-                    .unwrap_or_else(|| self.src(ret.range()).to_owned());
                 Some(format!("({args_str}) -> {ret_str}"))
             }
 
@@ -82,31 +108,53 @@ impl<'src> CallableReverse<'src> {
                 }
             }
 
-            Expr::Subscript(s) => {
-                let slice_rewrite = match s.slice.as_ref() {
-                    Expr::Tuple(t) if !t.parenthesized => {
-                        let rewrites: Vec<Option<String>> =
-                            t.elts.iter().map(|e| self.rewrite(e)).collect();
-                        if rewrites.iter().any(Option::is_some) {
-                            let parts: Vec<String> = rewrites
-                                .into_iter()
-                                .zip(t.elts.iter())
-                                .map(|(r, e)| r.unwrap_or_else(|| self.src(e.range()).to_owned()))
-                                .collect();
-                            Some(parts.join(", "))
-                        } else {
-                            None
-                        }
-                    }
-                    slice => self.rewrite(slice),
-                };
-                slice_rewrite.map(|s_text| format!("{}[{s_text}]", self.src(s.value.range())))
+            Expr::Subscript(s) => self.rewrite_subscript_children(s),
+
+            // descend into list literals (e.g. a `Callable[[A, B], R]` left
+            // intact in stub mode) so nested callable forms are still rewritten
+            Expr::List(l) => {
+                let rewrites: Vec<Option<String>> =
+                    l.elts.iter().map(|e| self.rewrite(e)).collect();
+                if rewrites.iter().any(Option::is_some) {
+                    let parts: Vec<String> = rewrites
+                        .into_iter()
+                        .zip(l.elts.iter())
+                        .map(|(r, e)| r.unwrap_or_else(|| self.src(e.range()).to_owned()))
+                        .collect();
+                    Some(format!("[{}]", parts.join(", ")))
+                } else {
+                    None
+                }
             }
 
             _ => None,
         }
     }
 
+    /// recurse into a subscript's slice, rewriting any nested callable forms
+    /// while keeping the `value[...]` wrapper. returns `None` if nothing in
+    /// the slice changed
+    fn rewrite_subscript_children(&mut self, s: &ExprSubscript) -> Option<String> {
+        let slice_rewrite = match s.slice.as_ref() {
+            Expr::Tuple(t) if !t.parenthesized => {
+                let rewrites: Vec<Option<String>> =
+                    t.elts.iter().map(|e| self.rewrite(e)).collect();
+                if rewrites.iter().any(Option::is_some) {
+                    let parts: Vec<String> = rewrites
+                        .into_iter()
+                        .zip(t.elts.iter())
+                        .map(|(r, e)| r.unwrap_or_else(|| self.src(e.range()).to_owned()))
+                        .collect();
+                    Some(parts.join(", "))
+                } else {
+                    None
+                }
+            }
+            slice => self.rewrite(slice),
+        };
+        slice_rewrite.map(|s_text| format!("{}[{s_text}]", self.src(s.value.range())))
+    }
+
     fn visit_annotation(&mut self, ann: &Expr) {
         if let Some(rewrite) = self.rewrite(ann) {
             self.edits.push(Fix::safe_edit(Edit::range_replacement(
@@ -137,6 +185,33 @@ mod tests {
         );
     }
 
+    fn check_stub(input: &str, expected: &str) {
+        let config = Config {
+            is_stub: true,
+            ..Config::test_default()
+        };
+        assert_eq!(reverse_transpile(input, &config).unwrap(), expected);
+    }
+
+    #[test]
+    fn stub_keeps_list_form_but_rewrites_ellipsis() {
+        // in stub mode the `Callable[[A], R]` list form is preserved (can't
+        // carry `Unpack[Ts]`/`*Ts` through the arrow), but the gradual
+        // `Callable[..., R]` form is still rewritten to `(...) -> R`
+        check_stub(
+            "from typing import Callable\na: Callable[..., int]\nb: Callable[[int], str]\n",
+            "from typing import Callable\na: (...) -> int\nb: Callable[[int], str]\n",
+        );
+    }
+
+    #[test]
+    fn stub_rewrites_nested_ellipsis_inside_list_form() {
+        check_stub(
+            "from typing import Callable\na: Callable[[Callable[..., int]], str]\n",
+            "from typing import Callable\na: Callable[[(...) -> int], str]\n",
+        );
+    }
+
     #[test]
     fn simple_callable() {
         check(
@@ -161,6 +236,22 @@ mod tests {
         );
     }
 
+    #[test]
+    fn ellipsis_args() {
+        check(
+            "from typing import Callable\na: Callable[..., int]\n",
+            "from typing import Callable\na: (...) -> int\n",
+        );
+    }
+
+    #[test]
+    fn ellipsis_args_nested_return() {
+        check(
+            "from typing import Callable\na: Callable[..., Callable[[int], str]]\n",
+            "from typing import Callable\na: (...) -> (int) -> str\n",
+        );
+    }
+
     #[test]
     fn callable_in_union() {
         check(

crates/by_transforms/src/reverse_transforms/empty_declarations.rs

@@ -20,12 +20,29 @@ use ruff_python_ast::{Expr, Stmt, StmtClassDef, StmtFunctionDef};
 use ruff_text_size::{Ranged, TextRange, TextSize};
 
 pub(crate) struct EmptyDeclarations {
+    /// when reversing a non-stub `.py`, an abstract method keeps its `: ...`
+    /// body: the forward pass maps a bodyless `abstract def` to `: raise
+    /// NotImplementedError`, so stripping the body would not round-trip. in a
+    /// stub the body is dropped — bodyless is the stub idiom and the forward
+    /// pass re-emits `: ...` there
+    is_stub: bool,
     pub(crate) edits: Vec<Fix>,
 }
 
 impl EmptyDeclarations {
-    pub(crate) fn new() -> Self {
-        Self { edits: Vec::new() }
+    pub(crate) fn new(is_stub: bool) -> Self {
+        Self {
+            is_stub,
+            edits: Vec::new(),
+        }
+    }
+
+    fn is_abstract(func: &StmtFunctionDef) -> bool {
+        func.decorator_list.iter().any(|d| match &d.expression {
+            Expr::Name(n) => n.id.as_str() == "abstractmethod",
+            Expr::Attribute(a) => a.attr.id.as_str() == "abstractmethod",
+            _ => false,
+        })
     }
 
     fn is_ellipsis_body(body: &[Stmt]) -> bool {
@@ -58,9 +75,22 @@ impl EmptyDeclarations {
     }
 
     fn process_function(&mut self, func: &StmtFunctionDef) {
-        // Decorated functions belong to specialized reverse passes (e.g.
-        // overload) that know how to handle the whole group atomically.
-        if !func.decorator_list.is_empty() {
+        // `@overload`-decorated functions belong to the overload reverse pass,
+        // which strips the decorator and the `: ...` body together. other
+        // decorators (`@property`, `@deprecated`, modifier-backed ones like
+        // `@abstractmethod`/`@final`) are fine to strip the body from — the
+        // decorator/modifier survives in front of the now-bodyless def
+        if func
+            .decorator_list
+            .iter()
+            .any(|d| matches!(&d.expression, Expr::Name(n) if n.id.as_str() == "overload"))
+        {
+            return;
+        }
+        // outside a stub, an abstract method's `: ...` body must survive: the
+        // forward pass turns a bodyless `abstract def` into `: raise
+        // NotImplementedError`, so dropping it here would not round-trip
+        if !self.is_stub && Self::is_abstract(func) {
             return;
         }
         if !Self::is_ellipsis_body(&func.body) {
@@ -251,6 +281,71 @@ mod tests {
         );
     }
 
+    #[test]
+    fn property_decorated_function_stripped() {
+        // non-`@overload` decorators don't defer to the overload pass; the
+        // `: ...` body is stripped and the decorator survives in front
+        check(
+            indoc! {"
+                class A:
+                    @property
+                    def x(self) -> int: ...
+            "},
+            indoc! {"
+                class A:
+                    @property
+                    def x(self) -> int
+            "},
+        );
+    }
+
+    #[test]
+    fn abstract_function_keeps_body_in_non_stub() {
+        // non-stub: `@abstractmethod` reverses to `abstract` but the `: ...`
+        // body is kept — a bodyless `abstract def` forward-maps to `: raise
+        // NotImplementedError`, so stripping would not round-trip
+        check(
+            indoc! {"
+                from abc import abstractmethod
+                class A:
+                    @abstractmethod
+                    def f(self) -> None: ...
+            "},
+            indoc! {"
+                from abc import abstractmethod
+                class A:
+                    abstract def f(self) -> None: ...
+            "},
+        );
+    }
+
+    #[test]
+    fn abstract_function_stripped_in_stub() {
+        // stub: bodyless is the idiom and the forward pass re-emits `: ...`
+        // for an abstract method in a stub, so the body is dropped here
+        let config = Config {
+            is_stub: true,
+            ..Config::test_default()
+        };
+        assert_eq!(
+            reverse_transpile(
+                indoc! {"
+                    from abc import abstractmethod
+                    class A:
+                        @abstractmethod
+                        def f(self) -> None: ...
+                "},
+                &config,
+            )
+            .unwrap(),
+            indoc! {"
+                from abc import abstractmethod
+                class A:
+                    abstract def f(self) -> None
+            "},
+        );
+    }
+
     #[test]
     fn decorated_function_left_to_overload_pass() {
         // @overload-decorated stubs are handled by the overload reverse pass;