Use find_keyword helper function (#5983)

## Summary Use `find_keyword` helper function instead of reimplementing it. ## Test Plan `cargo test`
astral-sh · Jul 22, 2023 · c7e4c58 · c7e4c58
1 parent 6ff566f
commit c7e4c58
Show file tree

Hide file tree

Showing 11 changed files with 29 additions and 101 deletions.
diff --git a/crates/ruff/src/rules/airflow/rules/task_variable_name.rs b/crates/ruff/src/rules/airflow/rules/task_variable_name.rs
@@ -3,6 +3,7 @@ use rustpython_parser::ast::{Expr, Ranged};
 
 use ruff_diagnostics::{Diagnostic, Violation};
 use ruff_macros::{derive_message_formats, violation};
+use ruff_python_ast::helpers::find_keyword;
 use rustpython_parser::ast::Constant;
 
 use crate::checkers::ast::Checker;
@@ -73,9 +74,7 @@ pub(crate) fn variable_name_task_id(
     }
 
     // If the call doesn't have a `task_id` keyword argument, we can't do anything.
-    let keyword = keywords
-        .iter()
-        .find(|keyword| keyword.arg.as_ref().map_or(false, |arg| arg == "task_id"))?;
+    let keyword = find_keyword(keywords, "task_id")?;
 
     // If the keyword argument is not a string, we can't do anything.
     let task_id = match &keyword.value {

diff --git a/crates/ruff/src/rules/flake8_bandit/rules/shell_injection.rs b/crates/ruff/src/rules/flake8_bandit/rules/shell_injection.rs
@@ -4,7 +4,7 @@ use rustpython_parser::ast::{self, Constant, Expr, Keyword, Ranged};
 
 use ruff_diagnostics::{Diagnostic, Violation};
 use ruff_macros::{derive_message_formats, violation};
-use ruff_python_ast::helpers::Truthiness;
+use ruff_python_ast::helpers::{find_keyword, Truthiness};
 use ruff_python_semantic::SemanticModel;
 
 use crate::{
@@ -272,13 +272,10 @@ fn find_shell_keyword<'a>(
     keywords: &'a [Keyword],
     semantic: &SemanticModel,
 ) -> Option<ShellKeyword<'a>> {
-    keywords
-        .iter()
-        .find(|keyword| keyword.arg.as_ref().map_or(false, |arg| arg == "shell"))
-        .map(|keyword| ShellKeyword {
-            truthiness: Truthiness::from_expr(&keyword.value, |id| semantic.is_builtin(id)),
-            keyword,
-        })
+    find_keyword(keywords, "shell").map(|keyword| ShellKeyword {
+        truthiness: Truthiness::from_expr(&keyword.value, |id| semantic.is_builtin(id)),
+        keyword,
+    })
 }
 
 /// Return `true` if the value provided to the `shell` call seems safe. This is based on Bandit's

diff --git a/crates/ruff/src/rules/flake8_django/rules/locals_in_render_function.rs b/crates/ruff/src/rules/flake8_django/rules/locals_in_render_function.rs
@@ -2,6 +2,7 @@ use rustpython_parser::ast::{self, Expr, Keyword, Ranged};
 
 use ruff_diagnostics::{Diagnostic, Violation};
 use ruff_macros::{derive_message_formats, violation};
+use ruff_python_ast::helpers::find_keyword;
 use ruff_python_semantic::SemanticModel;
 
 use crate::checkers::ast::Checker;
@@ -65,10 +66,7 @@ pub(crate) fn locals_in_render_function(
             return;
         }
         &args[2]
-    } else if let Some(keyword) = keywords
-        .iter()
-        .find(|keyword| keyword.arg.as_ref().map_or(false, |arg| arg == "context"))
-    {
+    } else if let Some(keyword) = find_keyword(keywords, "context") {
         if !is_locals_call(&keyword.value, checker.semantic()) {
             return;
         }

diff --git a/crates/ruff/src/rules/flake8_print/rules/print_call.rs b/crates/ruff/src/rules/flake8_print/rules/print_call.rs
@@ -2,7 +2,7 @@ use rustpython_parser::ast::{Expr, Keyword, Ranged};
 
 use ruff_diagnostics::{Diagnostic, Violation};
 use ruff_macros::{derive_message_formats, violation};
-use ruff_python_ast::helpers::is_const_none;
+use ruff_python_ast::helpers::{find_keyword, is_const_none};
 
 use crate::checkers::ast::Checker;
 use crate::registry::AsRule;
@@ -84,10 +84,7 @@ pub(crate) fn print_call(checker: &mut Checker, func: &Expr, keywords: &[Keyword
         }) {
             // If the print call has a `file=` argument (that isn't `None`, `"sys.stdout"`,
             // or `"sys.stderr"`), don't trigger T201.
-            if let Some(keyword) = keywords
-                .iter()
-                .find(|keyword| keyword.arg.as_ref().map_or(false, |arg| arg == "file"))
-            {
+            if let Some(keyword) = find_keyword(keywords, "file") {
                 if !is_const_none(&keyword.value) {
                     if checker.semantic().resolve_call_path(&keyword.value).map_or(
                         true,

diff --git a/crates/ruff/src/rules/pandas_vet/rules/read_table.rs b/crates/ruff/src/rules/pandas_vet/rules/read_table.rs
@@ -3,6 +3,7 @@ use rustpython_parser::ast::{Constant, Expr, Keyword, Ranged};
 
 use ruff_diagnostics::{Diagnostic, Violation};
 use ruff_macros::{derive_message_formats, violation};
+use ruff_python_ast::helpers::find_keyword;
 
 use crate::checkers::ast::Checker;
 
@@ -56,15 +57,7 @@ pub(crate) fn use_of_read_table(checker: &mut Checker, func: &Expr, keywords: &[
         if let Some(Expr::Constant(ast::ExprConstant {
             value: Constant::Str(value),
             ..
-        })) = keywords
-            .iter()
-            .find(|keyword| {
-                keyword
-                    .arg
-                    .as_ref()
-                    .map_or(false, |keyword| keyword.as_str() == "sep")
-            })
-            .map(|keyword| &keyword.value)
+        })) = find_keyword(keywords, "sep").map(|keyword| &keyword.value)
         {
             if value.as_str() == "," {
                 checker

diff --git a/crates/ruff/src/rules/pylint/helpers.rs b/crates/ruff/src/rules/pylint/helpers.rs
@@ -1,5 +1,6 @@
 use std::fmt;
 
+use ruff_python_ast::helpers::find_keyword;
 use rustpython_parser::ast;
 use rustpython_parser::ast::{CmpOp, Constant, Expr, Keyword};
 
@@ -11,14 +12,7 @@ use crate::settings::Settings;
 /// Returns the value of the `name` parameter to, e.g., a `TypeVar` constructor.
 pub(super) fn type_param_name<'a>(args: &'a [Expr], keywords: &'a [Keyword]) -> Option<&'a str> {
     // Handle both `TypeVar("T")` and `TypeVar(name="T")`.
-    let name_param = keywords
-        .iter()
-        .find(|keyword| {
-            keyword
-                .arg
-                .as_ref()
-                .map_or(false, |keyword| keyword.as_str() == "name")
-        })
+    let name_param = find_keyword(keywords, "name")
         .map(|keyword| &keyword.value)
         .or_else(|| args.get(0))?;
     if let Expr::Constant(ast::ExprConstant {

diff --git a/crates/ruff/src/rules/pylint/rules/invalid_envvar_default.rs b/crates/ruff/src/rules/pylint/rules/invalid_envvar_default.rs
@@ -2,6 +2,7 @@ use rustpython_parser::ast::{self, Constant, Expr, Keyword, Operator, Ranged};
 
 use ruff_diagnostics::{Diagnostic, Violation};
 use ruff_macros::{derive_message_formats, violation};
+use ruff_python_ast::helpers::find_keyword;
 
 use crate::checkers::ast::Checker;
 
@@ -92,14 +93,7 @@ pub(crate) fn invalid_envvar_default(
     {
         // Find the `default` argument, if it exists.
         let Some(expr) = args.get(1).or_else(|| {
-            keywords
-                .iter()
-                .find(|keyword| {
-                    keyword
-                        .arg
-                        .as_ref()
-                        .map_or(false, |arg| arg.as_str() == "default")
-                })
+            find_keyword(keywords, "default")
                 .map(|keyword| &keyword.value)
         }) else {
             return;

diff --git a/crates/ruff/src/rules/pylint/rules/invalid_envvar_value.rs b/crates/ruff/src/rules/pylint/rules/invalid_envvar_value.rs
@@ -2,6 +2,7 @@ use rustpython_parser::ast::{self, Constant, Expr, Keyword, Operator, Ranged};
 
 use ruff_diagnostics::{Diagnostic, Violation};
 use ruff_macros::{derive_message_formats, violation};
+use ruff_python_ast::helpers::find_keyword;
 
 use crate::checkers::ast::Checker;
 
@@ -89,9 +90,7 @@ pub(crate) fn invalid_envvar_value(
     {
         // Find the `key` argument, if it exists.
         let Some(expr) = args.get(0).or_else(|| {
-            keywords
-                .iter()
-                .find(|keyword| keyword.arg.as_ref().map_or(false, |arg| arg == "key"))
+            find_keyword(keywords, "key")
                 .map(|keyword| &keyword.value)
         }) else {
             return;

diff --git a/crates/ruff/src/rules/pylint/rules/type_bivariance.rs b/crates/ruff/src/rules/pylint/rules/type_bivariance.rs
@@ -4,7 +4,7 @@ use rustpython_parser::ast::{self, Expr, Ranged};
 
 use ruff_diagnostics::{Diagnostic, Violation};
 use ruff_macros::{derive_message_formats, violation};
-use ruff_python_ast::helpers::is_const_true;
+use ruff_python_ast::helpers::{find_keyword, is_const_true};
 
 use crate::checkers::ast::Checker;
 use crate::rules::pylint::helpers::type_param_name;
@@ -84,27 +84,13 @@ pub(crate) fn type_bivariance(checker: &mut Checker, value: &Expr) {
         return;
     };
 
-    let Some(covariant) = keywords
-        .iter()
-        .find(|keyword| {
-            keyword
-                .arg
-                .as_ref()
-                .map_or(false, |keyword| keyword.as_str() == "covariant")
-        })
+    let Some(covariant) = find_keyword(keywords, "covariant")
         .map(|keyword| &keyword.value)
     else {
         return;
     };
 
-    let Some(contravariant) = keywords
-        .iter()
-        .find(|keyword| {
-            keyword
-                .arg
-                .as_ref()
-                .map_or(false, |keyword| keyword.as_str() == "contravariant")
-        })
+    let Some(contravariant) = find_keyword(keywords, "contravariant")
         .map(|keyword| &keyword.value)
     else {
         return;

diff --git a/crates/ruff/src/rules/pylint/rules/type_name_incorrect_variance.rs b/crates/ruff/src/rules/pylint/rules/type_name_incorrect_variance.rs
@@ -4,7 +4,7 @@ use rustpython_parser::ast::{self, Expr, Ranged};
 
 use ruff_diagnostics::{Diagnostic, Violation};
 use ruff_macros::{derive_message_formats, violation};
-use ruff_python_ast::helpers::is_const_true;
+use ruff_python_ast::helpers::{find_keyword, is_const_true};
 
 use crate::checkers::ast::Checker;
 use crate::rules::pylint::helpers::type_param_name;
@@ -79,25 +79,9 @@ pub(crate) fn type_name_incorrect_variance(checker: &mut Checker, value: &Expr)
         return;
     };
 
-    let covariant = keywords
-        .iter()
-        .find(|keyword| {
-            keyword
-                .arg
-                .as_ref()
-                .map_or(false, |keyword| keyword.as_str() == "covariant")
-        })
-        .map(|keyword| &keyword.value);
-
-    let contravariant = keywords
-        .iter()
-        .find(|keyword| {
-            keyword
-                .arg
-                .as_ref()
-                .map_or(false, |keyword| keyword.as_str() == "contravariant")
-        })
-        .map(|keyword| &keyword.value);
+    let covariant = find_keyword(keywords, "covariant").map(|keyword| &keyword.value);
+
+    let contravariant = find_keyword(keywords, "contravariant").map(|keyword| &keyword.value);
 
     if !mismatch(param_name, covariant, contravariant) {
         return;
@@ -130,7 +114,7 @@ pub(crate) fn type_name_incorrect_variance(checker: &mut Checker, value: &Expr)
         .trim_end_matches("_co")
         .trim_end_matches("_contra");
     let replacement_name: String = match variance {
-        VarVariance::Bivariance => return, // Bivariate type are invalid, so ignore them for this rule.
+        VarVariance::Bivariance => return, // Bivariate types are invalid, so ignore them for this rule.
         VarVariance::Covariance => format!("{name_root}_co"),
         VarVariance::Contravariance => format!("{name_root}_contra"),
         VarVariance::Invariance => name_root.to_string(),

diff --git a/crates/ruff/src/rules/pyupgrade/rules/convert_typed_dict_functional_to_class.rs b/crates/ruff/src/rules/pyupgrade/rules/convert_typed_dict_functional_to_class.rs
@@ -7,7 +7,7 @@ use rustpython_parser::ast::{
 
 use ruff_diagnostics::{AutofixKind, Diagnostic, Edit, Fix, Violation};
 use ruff_macros::{derive_message_formats, violation};
-use ruff_python_ast::helpers::is_dunder;
+use ruff_python_ast::helpers::{find_keyword, is_dunder};
 use ruff_python_ast::source_code::Generator;
 use ruff_python_semantic::SemanticModel;
 use ruff_python_stdlib::identifiers::is_identifier;
@@ -200,19 +200,6 @@ fn properties_from_keywords(keywords: &[Keyword]) -> Result<Vec<Stmt>> {
         .collect()
 }
 
-// The only way to have the `total` keyword is to use the args version, like:
-// ```
-// TypedDict('name', {'a': int}, total=True)
-// ```
-fn match_total_from_only_keyword(keywords: &[Keyword]) -> Option<&Keyword> {
-    keywords.iter().find(|keyword| {
-        let Some(arg) = &keyword.arg else {
-            return false;
-        };
-        arg.as_str() == "total"
-    })
-}
-
 fn match_properties_and_total<'a>(
     args: &'a [Expr],
     keywords: &'a [Keyword],
@@ -223,7 +210,7 @@ fn match_properties_and_total<'a>(
     // MyType = TypedDict('MyType', {'a': int, 'b': str}, a=int, b=str)
     // ```
     if let Some(dict) = args.get(1) {
-        let total = match_total_from_only_keyword(keywords);
+        let total = find_keyword(keywords, "total");
         match dict {
             Expr::Dict(ast::ExprDict {
                 keys,