[ruff] itertools.batched() without explicit strict (RUF054)

crates/ruff_linter/resources/test/fixtures/ruff/RUF054.py

@@ -0,0 +1,52 @@
+from itertools import batched, count, cycle, repeat
+
+# Errors
+batched()
+batched(range(3), 1)
+batched("abc", 2)
+batched([i for i in range(42)], some_n)
+batched((foo for foo in cycle()))
+batched(batched([1, 2, 3], strict=True))
+
+# OK
+batched(range(3), 0, strict=True)
+batched(["a", "b"], count, strict=False)
+batched(("a", "b", "c"), zip(repeat()), strict=True)
+
+# OK (infinite iterators).
+batched(cycle("ABCDEF"), 3)
+batched(count(), qux + lorem)
+batched(repeat(1), ipsum // 19 @ 0x1)
+batched(repeat(1, None))
+batched(repeat(1, times=None))
+
+# Errors (limited iterators).
+batched(repeat(1, 1))
+batched(repeat(1, times=4))
+
+
+import itertools
+
+# Errors
+itertools.batched()
+itertools.batched(range(3), 1)
+itertools.batched("abc", 2)
+itertools.batched([i for i in range(42)], some_n)
+itertools.batched((foo for foo in cycle()))
+itertools.batched(itertools.batched([1, 2, 3], strict=True))
+
+# OK
+itertools.batched(range(3), 0, strict=True)
+itertools.batched(["a", "b"], count, strict=False)
+itertools.batched(("a", "b", "c"), zip(repeat()), strict=True)
+
+# OK (infinite iterators).
+itertools.batched(cycle("ABCDEF"), 3)
+itertools.batched(count(), qux + lorem)
+itertools.batched(repeat(1), ipsum // 19 @ 0x1)
+itertools.batched(repeat(1, None))
+itertools.batched(repeat(1, times=None))
+
+# Errors (limited iterators).
+itertools.batched(repeat(1, 1))
+itertools.batched(repeat(1, times=4))

crates/ruff_linter/src/checkers/ast/analyze/expression.rs

@@ -1055,6 +1055,9 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
             if checker.enabled(Rule::UnsafeMarkupUse) {
                 ruff::rules::unsafe_markup_call(checker, call);
             }
+            if checker.enabled(Rule::BatchedWithoutExplicitStrict) {
+                ruff::rules::batched_without_explicit_strict(checker, call);
+            }
         }
         Expr::Dict(dict) => {
             if checker.any_enabled(&[

crates/ruff_linter/src/codes.rs

@@ -971,6 +971,7 @@ pub fn code_to_rule(linter: Linter, code: &str) -> Option<(RuleGroup, Rule)> {
         (Ruff, "035") => (RuleGroup::Preview, rules::ruff::rules::UnsafeMarkupUse),
         (Ruff, "036") => (RuleGroup::Preview, rules::ruff::rules::NoneNotAtEndOfUnion),
         (Ruff, "038") => (RuleGroup::Preview, rules::ruff::rules::RedundantBoolLiteral),
+        (Ruff, "054") => (RuleGroup::Preview, rules::ruff::rules::BatchedWithoutExplicitStrict),
         (Ruff, "100") => (RuleGroup::Stable, rules::ruff::rules::UnusedNOQA),
         (Ruff, "101") => (RuleGroup::Stable, rules::ruff::rules::RedirectedNOQA),
 

crates/ruff_linter/src/rules/flake8_bugbear/rules/zip_without_explicit_strict.rs

@@ -17,7 +17,7 @@ use crate::fix::edits::add_argument;
 /// iterable. This can lead to subtle bugs.
 ///
 /// Pass `strict=True` to raise a `ValueError` if the iterables are of
-/// non-uniform length. Alternatively, if the iterables are deliberately
+/// non-uniform length. Alternatively, if the iterables are deliberately of
 /// different lengths, pass `strict=False` to make the intention explicit.
 ///
 /// ## Example
@@ -89,7 +89,7 @@ pub(crate) fn zip_without_explicit_strict(checker: &mut Checker, call: &ast::Exp
 
 /// Return `true` if the [`Expr`] appears to be an infinite iterator (e.g., a call to
 /// `itertools.cycle` or similar).
-fn is_infinite_iterator(arg: &Expr, semantic: &SemanticModel) -> bool {
+pub(crate) fn is_infinite_iterator(arg: &Expr, semantic: &SemanticModel) -> bool {
     let Expr::Call(ast::ExprCall {
         func,
         arguments: Arguments { args, keywords, .. },

crates/ruff_linter/src/rules/ruff/mod.rs

@@ -414,6 +414,25 @@ mod tests {
         Ok(())
     }
 
+    #[test_case(Rule::BatchedWithoutExplicitStrict, Path::new("RUF054.py"))]
+    fn ruf054(rule_code: Rule, path: &Path) -> Result<()> {
+        let snapshot = format!(
+            "preview__{}_{}",
+            rule_code.noqa_code(),
+            path.to_string_lossy()
+        );
+        let diagnostics = test_path(
+            Path::new("ruff").join(path).as_path(),
+            &settings::LinterSettings {
+                preview: PreviewMode::Enabled,
+                target_version: PythonVersion::Py313,
+                ..settings::LinterSettings::for_rule(rule_code)
+            },
+        )?;
+        assert_messages!(snapshot, diagnostics);
+        Ok(())
+    }
+
     #[test_case(Rule::UnsafeMarkupUse, Path::new("RUF035_extend_markup_names.py"))]
     #[test_case(Rule::UnsafeMarkupUse, Path::new("RUF035_skip_early_out.py"))]
     fn extend_allowed_callable(rule_code: Rule, path: &Path) -> Result<()> {

crates/ruff_linter/src/rules/ruff/rules/batched_without_explicit_strict.rs

@@ -0,0 +1,101 @@
+use crate::checkers::ast::Checker;
+use crate::fix::edits::add_argument;
+use crate::rules::flake8_bugbear::rules::is_infinite_iterator;
+use crate::settings::types::PythonVersion;
+use ruff_diagnostics::{AlwaysFixableViolation, Applicability, Diagnostic, Fix};
+use ruff_macros::{derive_message_formats, violation};
+use ruff_python_ast::{Arguments, ExprCall};
+
+/// ## What it does
+/// Checks for `itertools.batched` calls without an explicit `strict` parameter.
+///
+/// ## Why is this bad?
+/// By default, if the length of the iterable is not divisible by
+/// the second argument to `itertools.batched`, the last batch
+/// will be shorter than the rest.
+///
+/// Pass `strict=True` to raise a `ValueError` if the batches are of non-uniform length.
+/// Otherwise, pass `strict=False` to make the intention explicit.
+///
+/// ## Example
+/// ```python
+/// itertools.batched(iterable, n)
+/// ```
+///
+/// Use instead:
+/// ```python
+/// itertools.batched(iterable, n, strict=True)
+/// ```
+///
+/// ## Fix safety
+/// This rule's fix is marked as unsafe for `batched` calls that contain
+/// `**kwargs`, as adding a `strict` keyword argument to such a call may lead
+/// to a duplicate keyword argument error.
+///
+/// ## References
+/// - [Python documentation: `batched`](https://docs.python.org/3/library/itertools.html#batched)
+#[violation]
+pub struct BatchedWithoutExplicitStrict;
+
+impl AlwaysFixableViolation for BatchedWithoutExplicitStrict {
+    #[derive_message_formats]
+    fn message(&self) -> String {
+        "`batched()` without an explicit `strict=` parameter".to_string()
+    }
+
+    fn fix_title(&self) -> String {
+        "Add explicit value for parameter `strict=`".to_string()
+    }
+}
+
+/// RUF054
+pub(crate) fn batched_without_explicit_strict(checker: &mut Checker, call: &ExprCall) {
+    if checker.settings.target_version < PythonVersion::Py313 {
+        return;
+    }
+
+    let semantic = checker.semantic();
+    let (func, arguments) = (&call.func, &call.arguments);
+
+    let Some(qualified_name) = semantic.resolve_qualified_name(func) else {
+        return;
+    };
+    let first_positional = arguments.args.first();
+
+    if !matches!(qualified_name.segments(), ["itertools", "batched"])
+        || arguments.find_keyword("strict").is_some()
+        || first_positional.is_some_and(|it| is_infinite_iterator(it, semantic))
+    {
+        return;
+    }
+
+    let diagnostic = Diagnostic::new(BatchedWithoutExplicitStrict, call.range);
+    let fix = add_strict_fix(checker, arguments);
+
+    checker.diagnostics.push(diagnostic.with_fix(fix));
+}
+
+#[inline]
+fn add_strict_fix(checker: &Checker, arguments: &Arguments) -> Fix {
+    let edit = add_argument(
+        "strict=False",
+        arguments,
+        checker.comment_ranges(),
+        checker.locator().contents(),
+    );
+    let applicability = if has_kwargs(arguments) {
+        Applicability::Unsafe
+    } else {
+        Applicability::Safe
+    };
+
+    Fix::applicable_edit(edit, applicability)
+}
+
+#[inline]
+fn has_kwargs(arguments: &Arguments) -> bool {
+    arguments
+        .keywords
+        .iter()
+        .any(|keyword| keyword.arg.is_none())
+}

crates/ruff_linter/src/rules/ruff/rules/mod.rs

@@ -2,6 +2,7 @@ pub(crate) use ambiguous_unicode_character::*;
 pub(crate) use assert_with_print_message::*;
 pub(crate) use assignment_in_assert::*;
 pub(crate) use asyncio_dangling_task::*;
+pub(crate) use batched_without_explicit_strict::*;
 pub(crate) use collection_literal_concatenation::*;
 pub(crate) use decimal_from_float_literal::*;
 pub(crate) use default_factory_kwarg::*;
@@ -40,6 +41,7 @@ mod ambiguous_unicode_character;
 mod assert_with_print_message;
 mod assignment_in_assert;
 mod asyncio_dangling_task;
+mod batched_without_explicit_strict;
 mod collection_literal_concatenation;
 mod confusables;
 mod decimal_from_float_literal;