From 74ecdc73acd10642fbd685c0619b831d5eb424f8 Mon Sep 17 00:00:00 2001
From: "Colin J. Fuller" <cjfuller@gmail.com>
Date: Tue, 6 Sep 2022 20:18:46 -0400
Subject: [PATCH] Handle E731 in type-annotated assignment (#116)

---
 resources/test/fixtures/E731.py |  4 ++++
 src/check_ast.rs                | 17 ++++++++++++++++-
 src/linter.rs                   | 17 ++++++++++++-----
 3 files changed, 32 insertions(+), 6 deletions(-)

diff --git a/resources/test/fixtures/E731.py b/resources/test/fixtures/E731.py
index 8057f9cc07..13b151738b 100644
--- a/resources/test/fixtures/E731.py
+++ b/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))
diff --git a/src/check_ast.rs b/src/check_ast.rs
index 43f0fa04b8..b97e295c45 100644
--- a/src/check_ast.rs
+++ b/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;
             }
             _ => {}
diff --git a/src/linter.rs b/src/linter.rs
index cd731c9d17..1ef914ffca 100644
--- a/src/linter.rs
+++ b/src/linter.rs
@@ -268,11 +268,18 @@ mod tests {
             },
             &autofix::Mode::Generate,
         )?;
-        let expected = vec![Check {
-            kind: CheckKind::DoNotAssignLambda,
-            location: Location::new(1, 1),
-            fix: None,
-        }];
+        let expected = vec![
+            Check {
+                kind: CheckKind::DoNotAssignLambda,
+                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() {