From e6a44b80ca5f1020e8259acda01d11b3c00baeca Mon Sep 17 00:00:00 2001
From: Jonathan Plasse <13716151+JonathanPlasse@users.noreply.github.com>
Date: Sun, 23 Apr 2023 23:18:01 +0200
Subject: [PATCH] Add in_test to Context

---
 crates/ruff/src/checkers/ast/mod.rs           | 45 ++++++++++++++++++-
 .../src/analyze/typing.rs                     |  1 +
 crates/ruff_python_semantic/src/context.rs    |  2 +
 3 files changed, 46 insertions(+), 2 deletions(-)

diff --git a/crates/ruff/src/checkers/ast/mod.rs b/crates/ruff/src/checkers/ast/mod.rs
index 198c7c79f0d7c..2c84735a95802 100644
--- a/crates/ruff/src/checkers/ast/mod.rs
+++ b/crates/ruff/src/checkers/ast/mod.rs
@@ -151,6 +151,16 @@ macro_rules! visit_type_definition {
     }};
 }
 
+/// Visit an [`Expr`], and treat it as a test.
+macro_rules! visit_test {
+    ($self:ident, $expr:expr) => {{
+        let prev_in_test = $self.ctx.in_test;
+        $self.ctx.in_test = true;
+        $self.visit_expr($expr);
+        $self.ctx.in_test = prev_in_test;
+    }};
+}
+
 /// Visit an [`Expr`], and treat it as _not_ a type definition.
 macro_rules! visit_non_type_definition {
     ($self:ident, $expr:expr) => {{
@@ -2156,8 +2166,16 @@ where
                 }
                 self.visit_expr(target);
             }
+            StmtKind::Assert { test, .. } => {
+                visit_test!(self, test);
+            }
+            StmtKind::While { test, body, orelse } => {
+                visit_test!(self, test);
+                self.visit_body(body);
+                self.visit_body(orelse);
+            }
             StmtKind::If { test, body, orelse } => {
-                self.visit_expr(test);
+                visit_test!(self, test);
 
                 if flake8_type_checking::helpers::is_type_checking_block(&self.ctx, test) {
                     if self.settings.rules.enabled(Rule::EmptyTypeCheckingBlock) {
@@ -2244,6 +2262,11 @@ where
 
         let prev_in_literal = self.ctx.in_literal;
         let prev_in_type_definition = self.ctx.in_type_definition;
+        let prev_in_test = self.ctx.in_test;
+
+        if !matches!(expr.node, ExprKind::BoolOp { .. }) {
+            self.ctx.in_test = false;
+        }
 
         // Pre-visit.
         match &expr.node {
@@ -3584,6 +3607,11 @@ where
                     (self.ctx.scope_stack.clone(), self.ctx.parents.clone()),
                 ));
             }
+            ExprKind::IfExp { test, body, orelse } => {
+                visit_test!(self, test);
+                self.visit_expr(body);
+                self.visit_expr(orelse);
+            }
             ExprKind::Call {
                 func,
                 args,
@@ -3612,11 +3640,19 @@ where
                     .any(|target| call_path.as_slice() == ["mypy_extensions", target])
                     {
                         Some(Callable::MypyExtension)
+                    } else if call_path.as_slice() == ["", "bool"] && self.ctx.is_builtin("bool") {
+                        Some(Callable::Bool)
                     } else {
                         None
                     }
                 });
                 match callable {
+                    Some(Callable::Bool) => {
+                        self.visit_expr(func);
+                        if args.len() == 1 {
+                            visit_test!(self, &args[0]);
+                        }
+                    }
                     Some(Callable::Cast) => {
                         self.visit_expr(func);
                         if !args.is_empty() {
@@ -3815,6 +3851,7 @@ where
 
         self.ctx.in_type_definition = prev_in_type_definition;
         self.ctx.in_literal = prev_in_literal;
+        self.ctx.in_test = prev_in_test;
 
         self.ctx.pop_expr();
     }
@@ -3827,7 +3864,11 @@ where
                 &comprehension.iter,
             );
         }
-        visitor::walk_comprehension(self, comprehension);
+        self.visit_expr(&comprehension.iter);
+        self.visit_expr(&comprehension.target);
+        for expr in &comprehension.ifs {
+            visit_test!(self, expr);
+        }
     }
 
     fn visit_excepthandler(&mut self, excepthandler: &'b Excepthandler) {
diff --git a/crates/ruff_python_semantic/src/analyze/typing.rs b/crates/ruff_python_semantic/src/analyze/typing.rs
index 189b061e867e9..3c45adb9015bc 100644
--- a/crates/ruff_python_semantic/src/analyze/typing.rs
+++ b/crates/ruff_python_semantic/src/analyze/typing.rs
@@ -10,6 +10,7 @@ use crate::context::Context;
 
 #[derive(Copy, Clone)]
 pub enum Callable {
+    Bool,
     Cast,
     NewType,
     TypeVar,
diff --git a/crates/ruff_python_semantic/src/context.rs b/crates/ruff_python_semantic/src/context.rs
index 1d67601b52b71..ee5b0147bf875 100644
--- a/crates/ruff_python_semantic/src/context.rs
+++ b/crates/ruff_python_semantic/src/context.rs
@@ -48,6 +48,7 @@ pub struct Context<'a> {
     pub in_deferred_type_definition: bool,
     pub in_exception_handler: bool,
     pub in_f_string: bool,
+    pub in_test: bool,
     pub in_literal: bool,
     pub in_subscript: bool,
     pub in_type_checking_block: bool,
@@ -88,6 +89,7 @@ impl<'a> Context<'a> {
             in_deferred_type_definition: false,
             in_exception_handler: false,
             in_f_string: false,
+            in_test: false,
             in_literal: false,
             in_subscript: false,
             in_type_checking_block: false,