Preserve type annotations when fixing E731

crates/ruff/resources/test/fixtures/pycodestyle/E731.py

@@ -21,3 +21,30 @@ class F:
 f.append(lambda x: x**2)
 f = g = lambda x: x**2
 lambda: "no-op"
+
+# Annotated
+from typing import Callable, ParamSpec
+
+P = ParamSpec("P")
+
+# ParamSpec cannot be used in this context, so do not preserve the annotation.
+f: Callable[P, int] = lambda *args: len(args)
+f: Callable[[], None] = lambda: None
+f: Callable[..., None] = lambda a, b: None
+f: Callable[[int], int] = lambda x: 2 * x
+
+# Let's use the `Callable` type from `collections.abc` instead.
+from collections.abc import Callable
+
+f: Callable[[str, int], str] = lambda a, b: a * b
+f: Callable[[str, int], tuple[str, int]] = lambda a, b: (a, b)
+f: Callable[[str, int, list[str]], list[str]] = lambda a, b, /, c: [*c, a * b]
+
+
+# Override `Callable`
+class Callable:
+    pass
+
+
+# Do not copy the annotation from here on out.
+f: Callable[[str, int], str] = lambda a, b: a * b

crates/ruff/src/checkers/ast/mod.rs

@@ -1786,7 +1786,7 @@ where
             StmtKind::Assign { targets, value, .. } => {
                 if self.settings.rules.enabled(Rule::LambdaAssignment) {
                     if let [target] = &targets[..] {
-                        pycodestyle::rules::lambda_assignment(self, target, value, stmt);
+                        pycodestyle::rules::lambda_assignment(self, target, value, None, stmt);
                     }
                 }
 
@@ -1861,7 +1861,13 @@ where
             } => {
                 if self.settings.rules.enabled(Rule::LambdaAssignment) {
                     if let Some(value) = value {
-                        pycodestyle::rules::lambda_assignment(self, target, value, stmt);
+                        pycodestyle::rules::lambda_assignment(
+                            self,
+                            target,
+                            value,
+                            Some(annotation),
+                            stmt,
+                        );
                     }
                 }
                 if self

crates/ruff/src/rules/pycodestyle/rules/lambda_assignment.rs

@@ -1,4 +1,7 @@
-use rustpython_parser::ast::{Arguments, Expr, ExprKind, Location, Stmt, StmtKind};
+use ruff_python_semantic::context::Context;
+use rustpython_parser::ast::{
+    Arg, ArgData, Arguments, Constant, Expr, ExprKind, Location, Stmt, StmtKind,
+};
 
 use ruff_diagnostics::{AutofixKind, Diagnostic, Edit, Violation};
 use ruff_macros::{derive_message_formats, violation};
@@ -57,7 +60,13 @@ impl Violation for LambdaAssignment {
 }
 
 /// E731
-pub fn lambda_assignment(checker: &mut Checker, target: &Expr, value: &Expr, stmt: &Stmt) {
+pub fn lambda_assignment(
+    checker: &mut Checker,
+    target: &Expr,
+    value: &Expr,
+    annotation: Option<&Expr>,
+    stmt: &Stmt,
+) {
     if let ExprKind::Name { id, .. } = &target.node {
         if let ExprKind::Lambda { args, body } = &value.node {
             // If the assignment is in a class body, it might not be safe
@@ -87,9 +96,10 @@ pub fn lambda_assignment(checker: &mut Checker, target: &Expr, value: &Expr, stm
                 ));
                 let indentation = &leading_space(first_line);
                 let mut indented = String::new();
-                for (idx, line) in function(id, args, body, checker.stylist)
-                    .universal_newlines()
-                    .enumerate()
+                for (idx, line) in
+                    function(&checker.ctx, id, args, body, annotation, checker.stylist)
+                        .universal_newlines()
+                        .enumerate()
                 {
                     if idx == 0 {
                         indented.push_str(line);
@@ -111,14 +121,111 @@ pub fn lambda_assignment(checker: &mut Checker, target: &Expr, value: &Expr, stm
     }
 }
 
-fn function(name: &str, args: &Arguments, body: &Expr, stylist: &Stylist) -> String {
+/// Extract the argument types and return type from a `Callable` annotation.
+/// The `Callable` import can be from either `collections.abc` or `typing`.
+/// If an ellipsis is used for the argument types, an empty list is returned.
+/// The returned values are cloned, so they can be used as-is.
+fn extract_types(ctx: &Context, annotation: &Expr) -> Option<(Vec<Expr>, Expr)> {
+    let ExprKind::Subscript { value, slice, .. } = &annotation.node else {
+        return None;
+    };
+    let ExprKind::Tuple { elts, .. } = &slice.node else {
+        return None;
+    };
+    if elts.len() != 2 {
+        return None;
+    }
+    // The first argument to `Callable` must be a list of types, parameter
+    // specification, or ellipsis.
+    let args = match &elts[0].node {
+        ExprKind::List { elts, .. } => elts.clone(),
+        ExprKind::Constant {
+            value: Constant::Ellipsis,
+            ..
+        } => vec![],
+        _ => return None,
+    };
+    if !ctx.resolve_call_path(value).map_or(false, |call_path| {
+        call_path.as_slice() == ["collections", "abc", "Callable"]
+            || ctx.match_typing_call_path(&call_path, "Callable")
+    }) {
+        return None;
+    }
+    Some((args, elts[1].clone()))
+}
+
+fn function(
+    ctx: &Context,
+    name: &str,
+    args: &Arguments,
+    body: &Expr,
+    annotation: Option<&Expr>,
+    stylist: &Stylist,
+) -> String {
     let body = Stmt::new(
         Location::default(),
         Location::default(),
         StmtKind::Return {
             value: Some(Box::new(body.clone())),
         },
     );
+    if let Some(annotation) = annotation {
+        if let Some((arg_types, return_type)) = extract_types(ctx, annotation) {
+            // A `lambda` expression can only have positional and positional-only
+            // arguments. The order is always positional-only first, then positional.
+            let new_posonlyargs = args
+                .posonlyargs
+                .iter()
+                .enumerate()
+                .map(|(idx, arg)| {
+                    Arg::new(
+                        Location::default(),
+                        Location::default(),
+                        ArgData {
+                            annotation: arg_types
+                                .get(idx)
+                                .map(|arg_type| Box::new(arg_type.clone())),
+                            ..arg.node.clone()
+                        },
+                    )
+                })
+                .collect::<Vec<_>>();
+            let new_args = args
+                .args
+                .iter()
+                .enumerate()
+                .map(|(idx, arg)| {
+                    Arg::new(
+                        Location::default(),
+                        Location::default(),
+                        ArgData {
+                            annotation: arg_types
+                                .get(idx + new_posonlyargs.len())
+                                .map(|arg_type| Box::new(arg_type.clone())),
+                            ..arg.node.clone()
+                        },
+                    )
+                })
+                .collect::<Vec<_>>();
+            let func = Stmt::new(
+                Location::default(),
+                Location::default(),
+                StmtKind::FunctionDef {
+                    name: name.to_string(),
+                    args: Box::new(Arguments {
+                        posonlyargs: new_posonlyargs,
+                        args: new_args,
+                        ..args.clone()
+                    }),
+                    body: vec![body],
+                    decorator_list: vec![],
+                    returns: Some(Box::new(return_type)),
+                    type_comment: None,
+                },
+            );
+            return unparse_stmt(&func, stylist);
+        }
+    }
     let func = Stmt::new(
         Location::default(),
         Location::default(),

crates/ruff/src/rules/pycodestyle/snapshots/ruff__rules__pycodestyle__tests__E731_E731.py.snap

@@ -118,4 +118,168 @@ E731.py:14:5: E731 Do not assign a `lambda` expression, use a `def`
 18 | f = object()
    |
 
+E731.py:31:1: E731 [*] Do not assign a `lambda` expression, use a `def`
+   |
+31 | # ParamSpec cannot be used in this context, so do not preserve the annotation.
+32 | f: Callable[P, int] = lambda *args: len(args)
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ E731
+33 | f: Callable[[], None] = lambda: None
+34 | f: Callable[..., None] = lambda a, b: None
+   |
+   = help: Rewrite `f` as a `def`
+
+ℹ Suggested fix
+28 28 | P = ParamSpec("P")
+29 29 | 
+30 30 | # ParamSpec cannot be used in this context, so do not preserve the annotation.
+31    |-f: Callable[P, int] = lambda *args: len(args)
+   31 |+def f(*args):
+   32 |+    return len(args)
+32 33 | f: Callable[[], None] = lambda: None
+33 34 | f: Callable[..., None] = lambda a, b: None
+34 35 | f: Callable[[int], int] = lambda x: 2 * x
+
+E731.py:32:1: E731 [*] Do not assign a `lambda` expression, use a `def`
+   |
+32 | # ParamSpec cannot be used in this context, so do not preserve the annotation.
+33 | f: Callable[P, int] = lambda *args: len(args)
+34 | f: Callable[[], None] = lambda: None
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ E731
+35 | f: Callable[..., None] = lambda a, b: None
+36 | f: Callable[[int], int] = lambda x: 2 * x
+   |
+   = help: Rewrite `f` as a `def`
+
+ℹ Suggested fix
+29 29 | 
+30 30 | # ParamSpec cannot be used in this context, so do not preserve the annotation.
+31 31 | f: Callable[P, int] = lambda *args: len(args)
+32    |-f: Callable[[], None] = lambda: None
+   32 |+def f() -> None:
+   33 |+    return None
+33 34 | f: Callable[..., None] = lambda a, b: None
+34 35 | f: Callable[[int], int] = lambda x: 2 * x
+35 36 | 
+
+E731.py:33:1: E731 [*] Do not assign a `lambda` expression, use a `def`
+   |
+33 | f: Callable[P, int] = lambda *args: len(args)
+34 | f: Callable[[], None] = lambda: None
+35 | f: Callable[..., None] = lambda a, b: None
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ E731
+36 | f: Callable[[int], int] = lambda x: 2 * x
+   |
+   = help: Rewrite `f` as a `def`
+
+ℹ Suggested fix
+30 30 | # ParamSpec cannot be used in this context, so do not preserve the annotation.
+31 31 | f: Callable[P, int] = lambda *args: len(args)
+32 32 | f: Callable[[], None] = lambda: None
+33    |-f: Callable[..., None] = lambda a, b: None
+   33 |+def f(a, b) -> None:
+   34 |+    return None
+34 35 | f: Callable[[int], int] = lambda x: 2 * x
+35 36 | 
+36 37 | # Let's use the `Callable` type from `collections.abc` instead.
+
+E731.py:34:1: E731 [*] Do not assign a `lambda` expression, use a `def`
+   |
+34 | f: Callable[[], None] = lambda: None
+35 | f: Callable[..., None] = lambda a, b: None
+36 | f: Callable[[int], int] = lambda x: 2 * x
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ E731
+37 | 
+38 | # Let's use the `Callable` type from `collections.abc` instead.
+   |
+   = help: Rewrite `f` as a `def`
+
+ℹ Suggested fix
+31 31 | f: Callable[P, int] = lambda *args: len(args)
+32 32 | f: Callable[[], None] = lambda: None
+33 33 | f: Callable[..., None] = lambda a, b: None
+34    |-f: Callable[[int], int] = lambda x: 2 * x
+   34 |+def f(x: int) -> int:
+   35 |+    return 2 * x
+35 36 | 
+36 37 | # Let's use the `Callable` type from `collections.abc` instead.
+37 38 | from collections.abc import Callable
+
+E731.py:39:1: E731 [*] Do not assign a `lambda` expression, use a `def`
+   |
+39 | from collections.abc import Callable
+40 | 
+41 | f: Callable[[str, int], str] = lambda a, b: a * b
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ E731
+42 | f: Callable[[str, int], tuple[str, int]] = lambda a, b: (a, b)
+43 | f: Callable[[str, int, list[str]], list[str]] = lambda a, b, /, c: [*c, a * b]
+   |
+   = help: Rewrite `f` as a `def`
+
+ℹ Suggested fix
+36 36 | # Let's use the `Callable` type from `collections.abc` instead.
+37 37 | from collections.abc import Callable
+38 38 | 
+39    |-f: Callable[[str, int], str] = lambda a, b: a * b
+   39 |+def f(a: str, b: int) -> str:
+   40 |+    return a * b
+40 41 | f: Callable[[str, int], tuple[str, int]] = lambda a, b: (a, b)
+41 42 | f: Callable[[str, int, list[str]], list[str]] = lambda a, b, /, c: [*c, a * b]
+42 43 | 
+
+E731.py:40:1: E731 [*] Do not assign a `lambda` expression, use a `def`
+   |
+40 | f: Callable[[str, int], str] = lambda a, b: a * b
+41 | f: Callable[[str, int], tuple[str, int]] = lambda a, b: (a, b)
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ E731
+42 | f: Callable[[str, int, list[str]], list[str]] = lambda a, b, /, c: [*c, a * b]
+   |
+   = help: Rewrite `f` as a `def`
+
+ℹ Suggested fix
+37 37 | from collections.abc import Callable
+38 38 | 
+39 39 | f: Callable[[str, int], str] = lambda a, b: a * b
+40    |-f: Callable[[str, int], tuple[str, int]] = lambda a, b: (a, b)
+   40 |+def f(a: str, b: int) -> tuple[str, int]:
+   41 |+    return a, b
+41 42 | f: Callable[[str, int, list[str]], list[str]] = lambda a, b, /, c: [*c, a * b]
+42 43 | 
+43 44 | 
+
+E731.py:41:1: E731 [*] Do not assign a `lambda` expression, use a `def`
+   |
+41 | f: Callable[[str, int], str] = lambda a, b: a * b
+42 | f: Callable[[str, int], tuple[str, int]] = lambda a, b: (a, b)
+43 | f: Callable[[str, int, list[str]], list[str]] = lambda a, b, /, c: [*c, a * b]
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ E731
+   |
+   = help: Rewrite `f` as a `def`
+
+ℹ Suggested fix
+38 38 | 
+39 39 | f: Callable[[str, int], str] = lambda a, b: a * b
+40 40 | f: Callable[[str, int], tuple[str, int]] = lambda a, b: (a, b)
+41    |-f: Callable[[str, int, list[str]], list[str]] = lambda a, b, /, c: [*c, a * b]
+   41 |+def f(a: str, b: int, /, c: list[str]) -> list[str]:
+   42 |+    return [*c, a * b]
+42 43 | 
+43 44 | 
+44 45 | # Override `Callable`
+
+E731.py:50:1: E731 [*] Do not assign a `lambda` expression, use a `def`
+   |
+50 | # Do not copy the annotation from here on out.
+51 | f: Callable[[str, int], str] = lambda a, b: a * b
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ E731
+   |
+   = help: Rewrite `f` as a `def`
+
+ℹ Suggested fix
+47 47 | 
+48 48 | 
+49 49 | # Do not copy the annotation from here on out.
+50    |-f: Callable[[str, int], str] = lambda a, b: a * b
+   50 |+def f(a, b):
+   51 |+    return a * b
+