Skip to content

Commit dadad0e

Browse files
Remove some allocations in argument detection (#5481)
## Summary Drive-by PR to remove some allocations around argument name matching.
1 parent d2450c2 commit dadad0e

File tree

6 files changed

+97
-97
lines changed

6 files changed

+97
-97
lines changed

crates/ruff/src/checkers/ast/mod.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -72,7 +72,7 @@ pub(crate) struct Checker<'a> {
7272
deferred: Deferred<'a>,
7373
pub(crate) diagnostics: Vec<Diagnostic>,
7474
// Check-specific state.
75-
pub(crate) flake8_bugbear_seen: Vec<&'a Expr>,
75+
pub(crate) flake8_bugbear_seen: Vec<&'a ast::ExprName>,
7676
}
7777

7878
impl<'a> Checker<'a> {

crates/ruff/src/rules/flake8_bandit/rules/request_without_timeout.rs

Lines changed: 20 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -1,31 +1,28 @@
1-
use rustpython_parser::ast::{self, Constant, Expr, Keyword, Ranged};
1+
use rustpython_parser::ast::{Expr, Keyword, Ranged};
22

33
use ruff_diagnostics::{Diagnostic, Violation};
44
use ruff_macros::{derive_message_formats, violation};
5-
use ruff_python_ast::helpers::SimpleCallArgs;
5+
use ruff_python_ast::helpers::{is_const_none, SimpleCallArgs};
66

77
use crate::checkers::ast::Checker;
88

99
#[violation]
1010
pub struct RequestWithoutTimeout {
11-
pub timeout: Option<String>,
11+
implicit: bool,
1212
}
1313

1414
impl Violation for RequestWithoutTimeout {
1515
#[derive_message_formats]
1616
fn message(&self) -> String {
17-
let RequestWithoutTimeout { timeout } = self;
18-
match timeout {
19-
Some(value) => {
20-
format!("Probable use of requests call with timeout set to `{value}`")
21-
}
22-
None => format!("Probable use of requests call without timeout"),
17+
let RequestWithoutTimeout { implicit } = self;
18+
if *implicit {
19+
format!("Probable use of requests call without timeout")
20+
} else {
21+
format!("Probable use of requests call with timeout set to `None`")
2322
}
2423
}
2524
}
2625

27-
const HTTP_VERBS: [&str; 7] = ["get", "options", "head", "post", "put", "patch", "delete"];
28-
2926
/// S113
3027
pub(crate) fn request_without_timeout(
3128
checker: &mut Checker,
@@ -37,30 +34,26 @@ pub(crate) fn request_without_timeout(
3734
.semantic()
3835
.resolve_call_path(func)
3936
.map_or(false, |call_path| {
40-
HTTP_VERBS
41-
.iter()
42-
.any(|func_name| call_path.as_slice() == ["requests", func_name])
37+
matches!(
38+
call_path.as_slice(),
39+
[
40+
"requests",
41+
"get" | "options" | "head" | "post" | "put" | "patch" | "delete"
42+
]
43+
)
4344
})
4445
{
4546
let call_args = SimpleCallArgs::new(args, keywords);
46-
if let Some(timeout_arg) = call_args.keyword_argument("timeout") {
47-
if let Some(timeout) = match timeout_arg {
48-
Expr::Constant(ast::ExprConstant {
49-
value: value @ Constant::None,
50-
..
51-
}) => Some(checker.generator().constant(value)),
52-
_ => None,
53-
} {
47+
if let Some(timeout) = call_args.keyword_argument("timeout") {
48+
if is_const_none(timeout) {
5449
checker.diagnostics.push(Diagnostic::new(
55-
RequestWithoutTimeout {
56-
timeout: Some(timeout),
57-
},
58-
timeout_arg.range(),
50+
RequestWithoutTimeout { implicit: false },
51+
timeout.range(),
5952
));
6053
}
6154
} else {
6255
checker.diagnostics.push(Diagnostic::new(
63-
RequestWithoutTimeout { timeout: None },
56+
RequestWithoutTimeout { implicit: true },
6457
func.range(),
6558
));
6659
}

crates/ruff/src/rules/flake8_bugbear/rules/function_uses_loop_variable.rs

Lines changed: 45 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,8 @@
1-
use rustc_hash::FxHashSet;
21
use rustpython_parser::ast::{self, Comprehension, Expr, ExprContext, Ranged, Stmt};
32

43
use ruff_diagnostics::{Diagnostic, Violation};
54
use ruff_macros::{derive_message_formats, violation};
6-
use ruff_python_ast::helpers::collect_arg_names;
5+
use ruff_python_ast::helpers::includes_arg_name;
76
use ruff_python_ast::types::Node;
87
use ruff_python_ast::visitor;
98
use ruff_python_ast::visitor::Visitor;
@@ -58,19 +57,17 @@ impl Violation for FunctionUsesLoopVariable {
5857

5958
#[derive(Default)]
6059
struct LoadedNamesVisitor<'a> {
61-
// Tuple of: name, defining expression, and defining range.
62-
loaded: Vec<(&'a str, &'a Expr)>,
63-
// Tuple of: name, defining expression, and defining range.
64-
stored: Vec<(&'a str, &'a Expr)>,
60+
loaded: Vec<&'a ast::ExprName>,
61+
stored: Vec<&'a ast::ExprName>,
6562
}
6663

6764
/// `Visitor` to collect all used identifiers in a statement.
6865
impl<'a> Visitor<'a> for LoadedNamesVisitor<'a> {
6966
fn visit_expr(&mut self, expr: &'a Expr) {
7067
match expr {
71-
Expr::Name(ast::ExprName { id, ctx, range: _ }) => match ctx {
72-
ExprContext::Load => self.loaded.push((id, expr)),
73-
ExprContext::Store => self.stored.push((id, expr)),
68+
Expr::Name(name) => match &name.ctx {
69+
ExprContext::Load => self.loaded.push(name),
70+
ExprContext::Store => self.stored.push(name),
7471
ExprContext::Del => {}
7572
},
7673
_ => visitor::walk_expr(self, expr),
@@ -80,7 +77,7 @@ impl<'a> Visitor<'a> for LoadedNamesVisitor<'a> {
8077

8178
#[derive(Default)]
8279
struct SuspiciousVariablesVisitor<'a> {
83-
names: Vec<(&'a str, &'a Expr)>,
80+
names: Vec<&'a ast::ExprName>,
8481
safe_functions: Vec<&'a Expr>,
8582
}
8683

@@ -95,17 +92,20 @@ impl<'a> Visitor<'a> for SuspiciousVariablesVisitor<'a> {
9592
let mut visitor = LoadedNamesVisitor::default();
9693
visitor.visit_body(body);
9794

98-
// Collect all argument names.
99-
let mut arg_names = collect_arg_names(args);
100-
arg_names.extend(visitor.stored.iter().map(|(id, ..)| id));
101-
10295
// Treat any non-arguments as "suspicious".
103-
self.names.extend(
104-
visitor
105-
.loaded
106-
.into_iter()
107-
.filter(|(id, ..)| !arg_names.contains(id)),
108-
);
96+
self.names
97+
.extend(visitor.loaded.into_iter().filter(|loaded| {
98+
if visitor.stored.iter().any(|stored| stored.id == loaded.id) {
99+
return false;
100+
}
101+
102+
if includes_arg_name(&loaded.id, args) {
103+
return false;
104+
}
105+
106+
true
107+
}));
108+
109109
return;
110110
}
111111
Stmt::Return(ast::StmtReturn {
@@ -132,10 +132,9 @@ impl<'a> Visitor<'a> for SuspiciousVariablesVisitor<'a> {
132132
}) => {
133133
match func.as_ref() {
134134
Expr::Name(ast::ExprName { id, .. }) => {
135-
let id = id.as_str();
136-
if id == "filter" || id == "reduce" || id == "map" {
135+
if matches!(id.as_str(), "filter" | "reduce" | "map") {
137136
for arg in args {
138-
if matches!(arg, Expr::Lambda(_)) {
137+
if arg.is_lambda_expr() {
139138
self.safe_functions.push(arg);
140139
}
141140
}
@@ -159,7 +158,7 @@ impl<'a> Visitor<'a> for SuspiciousVariablesVisitor<'a> {
159158

160159
for keyword in keywords {
161160
if keyword.arg.as_ref().map_or(false, |arg| arg == "key")
162-
&& matches!(keyword.value, Expr::Lambda(_))
161+
&& keyword.value.is_lambda_expr()
163162
{
164163
self.safe_functions.push(&keyword.value);
165164
}
@@ -175,17 +174,19 @@ impl<'a> Visitor<'a> for SuspiciousVariablesVisitor<'a> {
175174
let mut visitor = LoadedNamesVisitor::default();
176175
visitor.visit_expr(body);
177176

178-
// Collect all argument names.
179-
let mut arg_names = collect_arg_names(args);
180-
arg_names.extend(visitor.stored.iter().map(|(id, ..)| id));
181-
182177
// Treat any non-arguments as "suspicious".
183-
self.names.extend(
184-
visitor
185-
.loaded
186-
.iter()
187-
.filter(|(id, ..)| !arg_names.contains(id)),
188-
);
178+
self.names
179+
.extend(visitor.loaded.into_iter().filter(|loaded| {
180+
if visitor.stored.iter().any(|stored| stored.id == loaded.id) {
181+
return false;
182+
}
183+
184+
if includes_arg_name(&loaded.id, args) {
185+
return false;
186+
}
187+
188+
true
189+
}));
189190

190191
return;
191192
}
@@ -198,15 +199,15 @@ impl<'a> Visitor<'a> for SuspiciousVariablesVisitor<'a> {
198199

199200
#[derive(Default)]
200201
struct NamesFromAssignmentsVisitor<'a> {
201-
names: FxHashSet<&'a str>,
202+
names: Vec<&'a str>,
202203
}
203204

204205
/// `Visitor` to collect all names used in an assignment expression.
205206
impl<'a> Visitor<'a> for NamesFromAssignmentsVisitor<'a> {
206207
fn visit_expr(&mut self, expr: &'a Expr) {
207208
match expr {
208209
Expr::Name(ast::ExprName { id, .. }) => {
209-
self.names.insert(id.as_str());
210+
self.names.push(id.as_str());
210211
}
211212
Expr::Starred(ast::ExprStarred { value, .. }) => {
212213
self.visit_expr(value);
@@ -223,7 +224,7 @@ impl<'a> Visitor<'a> for NamesFromAssignmentsVisitor<'a> {
223224

224225
#[derive(Default)]
225226
struct AssignedNamesVisitor<'a> {
226-
names: FxHashSet<&'a str>,
227+
names: Vec<&'a str>,
227228
}
228229

229230
/// `Visitor` to collect all used identifiers in a statement.
@@ -257,7 +258,7 @@ impl<'a> Visitor<'a> for AssignedNamesVisitor<'a> {
257258
}
258259

259260
fn visit_expr(&mut self, expr: &'a Expr) {
260-
if matches!(expr, Expr::Lambda(_)) {
261+
if expr.is_lambda_expr() {
261262
// Don't recurse.
262263
return;
263264
}
@@ -300,15 +301,15 @@ pub(crate) fn function_uses_loop_variable<'a>(checker: &mut Checker<'a>, node: &
300301

301302
// If a variable was used in a function or lambda body, and assigned in the
302303
// loop, flag it.
303-
for (name, expr) in suspicious_variables {
304-
if reassigned_in_loop.contains(name) {
305-
if !checker.flake8_bugbear_seen.contains(&expr) {
306-
checker.flake8_bugbear_seen.push(expr);
304+
for name in suspicious_variables {
305+
if reassigned_in_loop.contains(&name.id.as_str()) {
306+
if !checker.flake8_bugbear_seen.contains(&name) {
307+
checker.flake8_bugbear_seen.push(name);
307308
checker.diagnostics.push(Diagnostic::new(
308309
FunctionUsesLoopVariable {
309-
name: name.to_string(),
310+
name: name.id.to_string(),
310311
},
311-
expr.range(),
312+
name.range(),
312313
));
313314
}
314315
}

crates/ruff/src/rules/flake8_pytest_style/rules/fixture.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ use ruff_diagnostics::{AlwaysAutofixableViolation, Violation};
88
use ruff_diagnostics::{Diagnostic, Edit, Fix};
99
use ruff_macros::{derive_message_formats, violation};
1010
use ruff_python_ast::call_path::collect_call_path;
11-
use ruff_python_ast::helpers::collect_arg_names;
11+
use ruff_python_ast::helpers::includes_arg_name;
1212
use ruff_python_ast::identifier::Identifier;
1313
use ruff_python_ast::visitor;
1414
use ruff_python_ast::visitor::Visitor;
@@ -446,7 +446,7 @@ fn check_fixture_decorator_name(checker: &mut Checker, decorator: &Decorator) {
446446

447447
/// PT021
448448
fn check_fixture_addfinalizer(checker: &mut Checker, args: &Arguments, body: &[Stmt]) {
449-
if !collect_arg_names(args).contains(&"request") {
449+
if !includes_arg_name("request", args) {
450450
return;
451451
}
452452

crates/ruff/src/rules/flake8_pytest_style/rules/patch.rs

Lines changed: 11 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,9 @@
1-
use rustc_hash::FxHashSet;
2-
use rustpython_parser::ast::{self, Expr, Keyword, Ranged};
1+
use rustpython_parser::ast::{self, Arguments, Expr, Keyword, Ranged};
32

43
use ruff_diagnostics::{Diagnostic, Violation};
54
use ruff_macros::{derive_message_formats, violation};
65
use ruff_python_ast::call_path::collect_call_path;
7-
use ruff_python_ast::helpers::{collect_arg_names, SimpleCallArgs};
6+
use ruff_python_ast::helpers::{includes_arg_name, SimpleCallArgs};
87
use ruff_python_ast::visitor;
98
use ruff_python_ast::visitor::Visitor;
109

@@ -18,10 +17,10 @@ impl Violation for PytestPatchWithLambda {
1817
}
1918
}
2019

21-
#[derive(Default)]
2220
/// Visitor that checks references the argument names in the lambda body.
21+
#[derive(Debug)]
2322
struct LambdaBodyVisitor<'a> {
24-
names: FxHashSet<&'a str>,
23+
arguments: &'a Arguments,
2524
uses_args: bool,
2625
}
2726

@@ -32,11 +31,15 @@ where
3231
fn visit_expr(&mut self, expr: &'b Expr) {
3332
match expr {
3433
Expr::Name(ast::ExprName { id, .. }) => {
35-
if self.names.contains(&id.as_str()) {
34+
if includes_arg_name(id, self.arguments) {
3635
self.uses_args = true;
3736
}
3837
}
39-
_ => visitor::walk_expr(self, expr),
38+
_ => {
39+
if !self.uses_args {
40+
visitor::walk_expr(self, expr);
41+
}
42+
}
4043
}
4144
}
4245
}
@@ -60,7 +63,7 @@ fn check_patch_call(
6063
{
6164
// Walk the lambda body.
6265
let mut visitor = LambdaBodyVisitor {
63-
names: collect_arg_names(args),
66+
arguments: args,
6467
uses_args: false,
6568
};
6669
visitor.visit_expr(body);

0 commit comments

Comments
 (0)