Handle E731 in type-annotated assignment (#116)

resources/test/fixtures/E731.py

@@ -1,2 +1,6 @@
+from typing import Callable, Iterable
+
 a = lambda x: x**2
 b = map(lambda x: 2 * x, range(3))
+c: Callable = lambda x: x**2
+d: Iterable = map(lambda x: 2 * x, range(3))

src/check_ast.rs

@@ -391,7 +391,22 @@ impl Visitor for Checker<'_> {
                     }
                 }
             }
-            StmtKind::Delete { .. } | StmtKind::AnnAssign { .. } => {
+            StmtKind::AnnAssign { value, .. } => {
+                self.seen_non_import = true;
+                if self
+                    .settings
+                    .select
+                    .contains(CheckKind::DoNotAssignLambda.code())
+                {
+                    if let Some(v) = value {
+                        if let ExprKind::Lambda { .. } = v.node {
+                            self.checks
+                                .push(Check::new(CheckKind::DoNotAssignLambda, stmt.location));
+                        }
+                    }
+                }
+            }
+            StmtKind::Delete { .. } => {
                 self.seen_non_import = true;
             }
             _ => {}

src/linter.rs

@@ -268,11 +268,18 @@ mod tests {
             },
             &autofix::Mode::Generate,
         )?;
-        let expected = vec![Check {
+        let expected = vec![
+            Check {
                 kind: CheckKind::DoNotAssignLambda,
-            location: Location::new(1, 1),
+                location: Location::new(3, 1),
                 fix: None,
-        }];
+            },
+            Check {
+                kind: CheckKind::DoNotAssignLambda,
+                location: Location::new(5, 1),
+                fix: None,
+            },
+        ];
 
         assert_eq!(actual.len(), expected.len());
         for i in 0..actual.len() {