Skip to content

Commit

Permalink
Implement FLY002 (convert static string joins to f-strings)
Browse files Browse the repository at this point in the history
  • Loading branch information
akx committed May 4, 2023
1 parent 59d40f9 commit 129d55a
Show file tree
Hide file tree
Showing 11 changed files with 360 additions and 3 deletions.
18 changes: 18 additions & 0 deletions crates/ruff/resources/test/fixtures/flynt/FLY002.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,18 @@
import secrets
from random import random, choice

a = "Hello"
ok1 = " ".join([a, " World"]) # OK
ok2 = "".join(["Finally, ", a, " World"]) # OK
ok3 = "x".join(("1", "2", "3")) # OK
ok4 = "y".join([1, 2, 3]) # Technically OK, though would've been an error originally
ok5 = "a".join([random(), random()]) # OK (simple calls)
ok6 = "a".join([secrets.token_urlsafe(), secrets.token_hex()]) # OK (attr calls)

nok1 = "x".join({"4", "5", "yee"}) # Not OK (set)
nok2 = a.join(["1", "2", "3"]) # Not OK (not a static joiner)
nok3 = "a".join(a) # Not OK (not a static joinee)
nok4 = "a".join([a, a, *a]) # Not OK (not a static length)
nok5 = "a".join([choice("flarp")]) # Not OK (not a simple call)
nok6 = "a".join(x for x in "feefoofum") # Not OK (generator)
nok7 = "a".join([f"foo{8}", "bar"]) # Not OK (contains an f-string)
14 changes: 11 additions & 3 deletions crates/ruff/src/checkers/ast/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -43,9 +43,9 @@ use crate::rules::{
flake8_django, flake8_errmsg, flake8_gettext, flake8_implicit_str_concat,
flake8_import_conventions, flake8_logging_format, flake8_pie, flake8_print, flake8_pyi,
flake8_pytest_style, flake8_raise, flake8_return, flake8_self, flake8_simplify,
flake8_tidy_imports, flake8_type_checking, flake8_unused_arguments, flake8_use_pathlib, mccabe,
numpy, pandas_vet, pep8_naming, pycodestyle, pydocstyle, pyflakes, pygrep_hooks, pylint,
pyupgrade, ruff, tryceratops,
flake8_tidy_imports, flake8_type_checking, flake8_unused_arguments, flake8_use_pathlib, flynt,
mccabe, numpy, pandas_vet, pep8_naming, pycodestyle, pydocstyle, pyflakes, pygrep_hooks,
pylint, pyupgrade, ruff, tryceratops,
};
use crate::settings::types::PythonVersion;
use crate::settings::{flags, Settings};
Expand Down Expand Up @@ -2436,13 +2436,21 @@ where
// pyupgrade
Rule::FormatLiterals,
Rule::FString,
// flynt
Rule::StaticJoinToFString,
]) {
if let ExprKind::Attribute { value, attr, .. } = &func.node {
if let ExprKind::Constant {
value: Constant::Str(value),
..
} = &value.node
{
if attr == "join" {
// "...".join(...) call
if self.settings.rules.enabled(Rule::StaticJoinToFString) {
flynt::rules::static_join_to_fstring(self, expr, value);
}
}
if attr == "format" {
// "...".format(...) call
let location = expr.range();
Expand Down
4 changes: 4 additions & 0 deletions crates/ruff/src/codes.rs
Original file line number Diff line number Diff line change
Expand Up @@ -732,6 +732,10 @@ pub fn code_to_rule(linter: Linter, code: &str) -> Option<Rule> {
(Flake8Django, "012") => Rule::DjangoUnorderedBodyContentInModel,
(Flake8Django, "013") => Rule::DjangoNonLeadingReceiverDecorator,

// flynt
// Reserved: (Flynt, "001") => Rule::StringConcatenationToFString,
(Flynt, "002") => Rule::StaticJoinToFString,

_ => return None,
})
}
5 changes: 5 additions & 0 deletions crates/ruff/src/registry.rs
Original file line number Diff line number Diff line change
Expand Up @@ -666,6 +666,8 @@ ruff_macros::register_rules!(
rules::flake8_django::rules::DjangoModelWithoutDunderStr,
rules::flake8_django::rules::DjangoUnorderedBodyContentInModel,
rules::flake8_django::rules::DjangoNonLeadingReceiverDecorator,
// flynt
rules::flynt::rules::StaticJoinToFString,
);

pub trait AsRule {
Expand Down Expand Up @@ -821,6 +823,9 @@ pub enum Linter {
/// [tryceratops](https://pypi.org/project/tryceratops/1.1.0/)
#[prefix = "TRY"]
Tryceratops,
/// [flynt](https://pypi.org/project/flynt/)
#[prefix = "FLY"]
Flynt,
/// NumPy-specific rules
#[prefix = "NPY"]
Numpy,
Expand Down
64 changes: 64 additions & 0 deletions crates/ruff/src/rules/flynt/helpers.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,64 @@
use ruff_python_ast::helpers::create_expr;
use rustpython_parser::ast::{Constant, Expr, ExprKind};

fn to_formatted_value_expr(inner: Expr) -> Expr {
create_expr(ExprKind::FormattedValue {
value: Box::new(inner),
conversion: 0,
format_spec: None,
})
}

/// Figure out if `expr` represents a "simple" call
/// (i.e. one that can be safely converted to a formatted value).
fn is_simple_call(expr: &Expr) -> bool {
match &expr.node {
ExprKind::Call {
func,
args,
keywords,
} => args.is_empty() && keywords.is_empty() && is_simple_callee(func),
_ => false,
}
}

/// Figure out if `func` represents a "simple" callee (a bare name, or a chain of simple
/// attribute accesses).
fn is_simple_callee(func: &Expr) -> bool {
match &func.node {
ExprKind::Name { .. } => true,
ExprKind::Attribute { value, .. } => is_simple_callee(value),
_ => false,
}
}

/// Convert an expression to a f-string element (if it looks like a good idea).
pub fn to_fstring_elem(expr: Expr) -> Option<Expr> {
match &expr.node {
// These are directly handled by `unparse_fstring_elem`:
ExprKind::Constant {
value: Constant::Str(_),
..
}
| ExprKind::JoinedStr { .. }
| ExprKind::FormattedValue { .. } => Some(expr),
// These should be pretty safe to wrap in a formatted value.
ExprKind::Constant {
value:
Constant::Int(_) | Constant::Float(_) | Constant::Bool(_) | Constant::Complex { .. },
..
}
| ExprKind::Name { .. }
| ExprKind::Attribute { .. } => Some(to_formatted_value_expr(expr)),
ExprKind::Call { .. } if is_simple_call(&expr) => Some(to_formatted_value_expr(expr)),
_ => None,
}
}

/// Convert a string to a constant string expression.
pub fn to_constant_string(s: &str) -> Expr {
create_expr(ExprKind::Constant {
value: Constant::Str(s.to_owned()),
kind: None,
})
}
24 changes: 24 additions & 0 deletions crates/ruff/src/rules/flynt/mod.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,24 @@
//! Rules from [flynt](https://pypi.org/project/flynt/).
mod helpers;
pub(crate) mod rules;

#[cfg(test)]
mod tests {
use crate::registry::Rule;
use crate::test::test_path;
use crate::{assert_messages, settings};
use anyhow::Result;
use std::path::Path;
use test_case::test_case;

#[test_case(Rule::StaticJoinToFString, Path::new("FLY002.py"); "FLY002")]
fn rules(rule_code: Rule, path: &Path) -> Result<()> {
let snapshot = format!("{}_{}", rule_code.noqa_code(), path.to_string_lossy());
let diagnostics = test_path(
Path::new("flynt").join(path).as_path(),
&settings::Settings::for_rule(rule_code),
)?;
assert_messages!(snapshot, diagnostics);
Ok(())
}
}
3 changes: 3 additions & 0 deletions crates/ruff/src/rules/flynt/rules/mod.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
mod static_join_to_fstring;

pub use static_join_to_fstring::{static_join_to_fstring, StaticJoinToFString};
98 changes: 98 additions & 0 deletions crates/ruff/src/rules/flynt/rules/static_join_to_fstring.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,98 @@
use rustpython_parser::ast::{Expr, ExprKind};

use ruff_diagnostics::{AutofixKind, Diagnostic, Edit, Violation};
use ruff_macros::{derive_message_formats, violation};
use ruff_python_ast::helpers::{create_expr, unparse_expr};

use crate::checkers::ast::Checker;
use crate::registry::AsRule;
use crate::rules::flynt::helpers;

#[violation]
pub struct StaticJoinToFString {
pub expr: String,
pub fixable: bool,
}

impl Violation for StaticJoinToFString {
const AUTOFIX: AutofixKind = AutofixKind::Sometimes;

#[derive_message_formats]
fn message(&self) -> String {
let StaticJoinToFString { expr, .. } = self;
format!("Consider `{expr}` instead of string join")
}

fn autofix_title_formatter(&self) -> Option<fn(&Self) -> String> {
self.fixable
.then_some(|StaticJoinToFString { expr, .. }| format!("Replace with `{expr}`"))
}
}

fn is_static_length(elts: &[Expr]) -> bool {
elts.iter()
.all(|e| !matches!(e.node, ExprKind::Starred { .. }))
}

fn build_fstring(joiner: &str, joinees: &Vec<Expr>) -> Option<Expr> {
let mut fstring_elems = Vec::with_capacity(joinees.len() * 2);
for (i, expr) in joinees.iter().enumerate() {
if matches!(expr.node, ExprKind::JoinedStr { .. }) {
// Oops, already an f-string. We don't know how to handle those
// gracefully right now.
return None;
}
let elem = helpers::to_fstring_elem(expr.clone())?;
if i != 0 {
fstring_elems.push(helpers::to_constant_string(joiner));
}
fstring_elems.push(elem);
}
Some(create_expr(ExprKind::JoinedStr {
values: fstring_elems,
}))
}

pub fn static_join_to_fstring(checker: &mut Checker, expr: &Expr, joiner: &str) {
let ExprKind::Call {
func: _,
args,
keywords,
} = &expr.node else {
return;
};

if !keywords.is_empty() || args.len() != 1 {
// If there are kwargs or more than one argument,
// this is some non-standard string join call.
return;
}

// Get the elements to join; skip e.g. generators, sets, etc.
let joinees = match &args[0].node {
ExprKind::List { elts, .. } if is_static_length(elts) => elts,
ExprKind::Tuple { elts, .. } if is_static_length(elts) => elts,
_ => return,
};

// Try to build the fstring (internally checks whether e.g. the elements are
// convertible to f-string parts).
let Some(new_expr) = build_fstring(joiner, joinees) else { return };

let contents = unparse_expr(&new_expr, checker.stylist);
let fixable = true; // I'm not sure there is a case where this is not fixable..?

let mut diagnostic = Diagnostic::new(
StaticJoinToFString {
expr: contents.clone(),
fixable,
},
expr.range(),
);
if checker.patch(diagnostic.kind.rule()) {
if fixable {
diagnostic.set_fix(Edit::range_replacement(contents, expr.range()));
}
}
checker.diagnostics.push(diagnostic);
}
Loading

0 comments on commit 129d55a

Please sign in to comment.