Skip to content

Commit

Permalink
[flake8-simplify] Implement SIM911 (#9460)
Browse files Browse the repository at this point in the history
## Summary

Closes #9319, implements the [`SIM911` rule from
`flake8-simplify`](MartinThoma/flake8-simplify#183).


#### Note
I wasn't sure whether or not to include
```rs
if checker.settings.preview.is_disabled() {
    return;
}
```
at the beginning of the function with violation logic if the rule's
already declared as part of `RuleGroup::Preview`.
I've seen both variants, so I'd appreciate some feedback on that :)
  • Loading branch information
trag1c authored Jan 11, 2024
1 parent f192c72 commit eb4ed24
Show file tree
Hide file tree
Showing 6 changed files with 160 additions and 0 deletions.
Original file line number Diff line number Diff line change
@@ -0,0 +1,23 @@
def foo(d: dict[str, str]) -> None:
for k, v in zip(d.keys(), d.values()): # SIM911
...

for k, v in zip(d.keys(), d.values(), strict=True): # SIM911
...

for k, v in zip(d.keys(), d.values(), struct=True): # OK
...


d1 = d2 = {}

for k, v in zip(d1.keys(), d2.values()): # OK
...

for k, v in zip(d1.items(), d2.values()): # OK
...

for k, v in zip(d2.keys(), d2.values()): # SIM911
...

items = zip(x.keys(), x.values()) # OK
3 changes: 3 additions & 0 deletions crates/ruff_linter/src/checkers/ast/analyze/expression.rs
Original file line number Diff line number Diff line change
Expand Up @@ -863,6 +863,9 @@ pub(crate) fn expression(expr: &Expr, checker: &mut Checker) {
if checker.enabled(Rule::DictGetWithNoneDefault) {
flake8_simplify::rules::dict_get_with_none_default(checker, expr);
}
if checker.enabled(Rule::ZipDictKeysAndValues) {
flake8_simplify::rules::zip_dict_keys_and_values(checker, call);
}
if checker.any_enabled(&[
Rule::OsPathAbspath,
Rule::OsChmod,
Expand Down
1 change: 1 addition & 0 deletions crates/ruff_linter/src/codes.rs
Original file line number Diff line number Diff line change
Expand Up @@ -472,6 +472,7 @@ pub fn code_to_rule(linter: Linter, code: &str) -> Option<(RuleGroup, Rule)> {
(Flake8Simplify, "300") => (RuleGroup::Stable, rules::flake8_simplify::rules::YodaConditions),
(Flake8Simplify, "401") => (RuleGroup::Stable, rules::flake8_simplify::rules::IfElseBlockInsteadOfDictGet),
(Flake8Simplify, "910") => (RuleGroup::Stable, rules::flake8_simplify::rules::DictGetWithNoneDefault),
(Flake8Simplify, "911") => (RuleGroup::Preview, rules::flake8_simplify::rules::ZipDictKeysAndValues),

// flake8-copyright
#[allow(deprecated)]
Expand Down
2 changes: 2 additions & 0 deletions crates/ruff_linter/src/rules/flake8_simplify/rules/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ pub(crate) use reimplemented_builtin::*;
pub(crate) use return_in_try_except_finally::*;
pub(crate) use suppressible_exception::*;
pub(crate) use yoda_conditions::*;
pub(crate) use zip_dict_keys_and_values::*;

mod ast_bool_op;
mod ast_expr;
Expand All @@ -34,3 +35,4 @@ mod reimplemented_builtin;
mod return_in_try_except_finally;
mod suppressible_exception;
mod yoda_conditions;
mod zip_dict_keys_and_values;
Original file line number Diff line number Diff line change
@@ -0,0 +1,130 @@
use ast::{ExprAttribute, ExprName, Identifier};
use ruff_macros::{derive_message_formats, violation};
use ruff_python_ast::{self as ast, Arguments, Expr, ExprCall};
use ruff_text_size::Ranged;

use crate::{checkers::ast::Checker, fix::snippet::SourceCodeSnippet};
use ruff_diagnostics::{AlwaysFixableViolation, Diagnostic, Edit, Fix};
use ruff_python_semantic::analyze::typing::is_dict;

/// ## What it does
/// Checks for use of `zip()` to iterate over keys and values of a dictionary at once.
///
/// ## Why is this bad?
/// The `dict` type provides an `.items()` method which is faster and more readable.
///
/// ## Example
/// ```python
/// flag_stars = {"USA": 50, "Slovenia": 3, "Panama": 2, "Australia": 6}
///
/// for country, stars in zip(flag_stars.keys(), flag_stars.values()):
/// print(f"{country}'s flag has {stars} stars.")
/// ```
///
/// Use instead:
/// ```python
/// flag_stars = {"USA": 50, "Slovenia": 3, "Panama": 2, "Australia": 6}
///
/// for country, stars in flag_stars.items():
/// print(f"{country}'s flag has {stars} stars.")
/// ```
///
/// ## References
/// - [Python documentation: `dict.items`](https://docs.python.org/3/library/stdtypes.html#dict.items)
#[violation]
pub struct ZipDictKeysAndValues {
expected: SourceCodeSnippet,
actual: SourceCodeSnippet,
}

impl AlwaysFixableViolation for ZipDictKeysAndValues {
#[derive_message_formats]
fn message(&self) -> String {
let ZipDictKeysAndValues { expected, actual } = self;
if let (Some(expected), Some(actual)) = (expected.full_display(), actual.full_display()) {
format!("Use `{expected}` instead of `{actual}`")
} else {
format!("Use `dict.items()` instead of `zip(dict.keys(), dict.values())`")
}
}

fn fix_title(&self) -> String {
let ZipDictKeysAndValues { expected, actual } = self;
if let (Some(expected), Some(actual)) = (expected.full_display(), actual.full_display()) {
format!("Replace `{actual}` with `{expected}`")
} else {
"Replace `zip(dict.keys(), dict.values())` with `dict.items()`".to_string()
}
}
}

/// SIM911
pub(crate) fn zip_dict_keys_and_values(checker: &mut Checker, expr: &ExprCall) {
let ExprCall {
func,
arguments: Arguments { args, keywords, .. },
..
} = expr;
match &keywords[..] {
[] => {}
[ast::Keyword {
arg: Some(name), ..
}] if name.as_str() == "strict" => {}
_ => return,
};
if matches!(func.as_ref(), Expr::Name(ExprName { id, .. }) if id != "zip") {
return;
}
let [arg1, arg2] = &args[..] else {
return;
};
let Some((var1, attr1)) = get_var_attr(arg1) else {
return;
};
let Some((var2, attr2)) = get_var_attr(arg2) else {
return;
};
if var1.id != var2.id || attr1 != "keys" || attr2 != "values" {
return;
}

let Some(binding) = checker
.semantic()
.only_binding(var1)
.map(|id| checker.semantic().binding(id))
else {
return;
};
if !is_dict(binding, checker.semantic()) {
return;
}

let expected = format!("{}.items()", checker.locator().slice(var1));
let actual = checker.locator().slice(expr);

let mut diagnostic = Diagnostic::new(
ZipDictKeysAndValues {
expected: SourceCodeSnippet::new(expected.clone()),
actual: SourceCodeSnippet::from_str(actual),
},
expr.range(),
);
diagnostic.set_fix(Fix::safe_edit(Edit::range_replacement(
expected,
expr.range(),
)));
checker.diagnostics.push(diagnostic);
}

fn get_var_attr(expr: &Expr) -> Option<(&ExprName, &Identifier)> {
let Expr::Call(ast::ExprCall { func, .. }) = expr else {
return None;
};
let Expr::Attribute(ExprAttribute { value, attr, .. }) = func.as_ref() else {
return None;
};
let Expr::Name(var_name) = value.as_ref() else {
return None;
};
Some((var_name, attr))
}
1 change: 1 addition & 0 deletions ruff.schema.json

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

0 comments on commit eb4ed24

Please sign in to comment.