From 81de3a16bca69cdb23ef1f15e7d06f1be12c68af Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Tue, 28 Mar 2023 11:20:22 -0400 Subject: [PATCH] Include `with` statements in complexity calculation (#3771) --- crates/ruff/src/rules/mccabe/rules.rs | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/crates/ruff/src/rules/mccabe/rules.rs b/crates/ruff/src/rules/mccabe/rules.rs index 5b07789db6..db0e8a3f03 100644 --- a/crates/ruff/src/rules/mccabe/rules.rs +++ b/crates/ruff/src/rules/mccabe/rules.rs @@ -79,6 +79,9 @@ fn get_complexity_number(stmts: &[Stmt]) -> usize { complexity += get_complexity_number(body); complexity += get_complexity_number(orelse); } + StmtKind::With { body, .. } | StmtKind::AsyncWith { body, .. } => { + complexity += get_complexity_number(body); + } StmtKind::While { body, orelse, .. } => { complexity += 1; complexity += get_complexity_number(body); @@ -403,6 +406,19 @@ def process_detect_lines(): finally: if res: errors.append(f"Non-zero exit code {res}") +"#; + let stmts = parser::parse_program(source, "")?; + assert_eq!(get_complexity_number(&stmts), 2); + Ok(()) + } + + #[test] + fn with() -> Result<()> { + let source = r#" +def with_lock(): + with lock: + if foo: + print('bar') "#; let stmts = parser::parse_program(source, "")?; assert_eq!(get_complexity_number(&stmts), 2);