Add in_boolean_test to Context (#4072)

JonathanPlasse · web-flow · commit 5e91211e6d1b · 2023-04-23T23:18:23.000-06:00
diff --git a/crates/ruff/src/checkers/ast/mod.rs b/crates/ruff/src/checkers/ast/mod.rs
@@ -161,6 +161,18 @@ macro_rules! visit_non_type_definition {
     }};
 }
 
+/// Visit an [`Expr`], and treat it as a boolean test. This is useful for detecting whether an
+/// expressions return value is significant, or whether the calling context only relies on
+/// its truthiness.
+macro_rules! visit_boolean_test {
+    ($self:ident, $expr:expr) => {{
+        let prev_in_boolean_test = $self.ctx.in_boolean_test;
+        $self.ctx.in_boolean_test = true;
+        $self.visit_expr($expr);
+        $self.ctx.in_boolean_test = prev_in_boolean_test;
+    }};
+}
+
 impl<'a, 'b> Visitor<'b> for Checker<'a>
 where
     'b: 'a,
@@ -2154,8 +2166,19 @@ where
                 }
                 self.visit_expr(target);
             }
+            StmtKind::Assert { test, msg } => {
+                visit_boolean_test!(self, test);
+                if let Some(expr) = msg {
+                    self.visit_expr(expr);
+                }
+            }
+            StmtKind::While { test, body, orelse } => {
+                visit_boolean_test!(self, test);
+                self.visit_body(body);
+                self.visit_body(orelse);
+            }
             StmtKind::If { test, body, orelse } => {
-                self.visit_expr(test);
+                visit_boolean_test!(self, test);
 
                 if flake8_type_checking::helpers::is_type_checking_block(&self.ctx, test) {
                     if self.settings.rules.enabled(Rule::EmptyTypeCheckingBlock) {
@@ -2242,6 +2265,11 @@ where
 
         let prev_in_literal = self.ctx.in_literal;
         let prev_in_type_definition = self.ctx.in_type_definition;
+        let prev_in_boolean_test = self.ctx.in_boolean_test;
+
+        if !matches!(expr.node, ExprKind::BoolOp { .. }) {
+            self.ctx.in_boolean_test = false;
+        }
 
         // Pre-visit.
         match &expr.node {
@@ -3582,6 +3610,11 @@ where
                     (self.ctx.scope_stack.clone(), self.ctx.parents.clone()),
                 ));
             }
+            ExprKind::IfExp { test, body, orelse } => {
+                visit_boolean_test!(self, test);
+                self.visit_expr(body);
+                self.visit_expr(orelse);
+            }
             ExprKind::Call {
                 func,
                 args,
@@ -3610,11 +3643,22 @@ where
                     .any(|target| call_path.as_slice() == ["mypy_extensions", target])
                     {
                         Some(Callable::MypyExtension)
+                    } else if call_path.as_slice() == ["", "bool"] {
+                        Some(Callable::Bool)
                     } else {
                         None
                     }
                 });
                 match callable {
+                    Some(Callable::Bool) => {
+                        self.visit_expr(func);
+                        if !args.is_empty() {
+                            visit_boolean_test!(self, &args[0]);
+                        }
+                        for expr in args.iter().skip(1) {
+                            self.visit_expr(expr);
+                        }
+                    }
                     Some(Callable::Cast) => {
                         self.visit_expr(func);
                         if !args.is_empty() {
@@ -3813,6 +3857,7 @@ where
 
         self.ctx.in_type_definition = prev_in_type_definition;
         self.ctx.in_literal = prev_in_literal;
+        self.ctx.in_boolean_test = prev_in_boolean_test;
 
         self.ctx.pop_expr();
     }
@@ -3825,7 +3870,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_boolean_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
@@ -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
@@ -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_boolean_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_boolean_test: false,
             in_literal: false,
             in_subscript: false,
             in_type_checking_block: false,
