From f6e49cfff68cfc09be2b25cf74066ad9ea98fdc7 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Tue, 20 Feb 2024 13:02:01 -0500 Subject: [PATCH] Allow os.environ modifications between imports --- .../test/fixtures/pycodestyle/E402_2.py | 6 +++ crates/ruff_linter/src/checkers/ast/mod.rs | 4 +- .../ruff_linter/src/rules/pycodestyle/mod.rs | 2 + .../rules/module_import_not_at_top_of_file.rs | 5 +++ ...s__pycodestyle__tests__E402_E402_2.py.snap | 12 +++++ ...style__tests__preview__E402_E402_2.py.snap | 4 ++ .../src/analyze/imports.rs | 45 +++++++++++++++++++ 7 files changed, 77 insertions(+), 1 deletion(-) create mode 100644 crates/ruff_linter/resources/test/fixtures/pycodestyle/E402_2.py create mode 100644 crates/ruff_linter/src/rules/pycodestyle/snapshots/ruff_linter__rules__pycodestyle__tests__E402_E402_2.py.snap create mode 100644 crates/ruff_linter/src/rules/pycodestyle/snapshots/ruff_linter__rules__pycodestyle__tests__preview__E402_E402_2.py.snap diff --git a/crates/ruff_linter/resources/test/fixtures/pycodestyle/E402_2.py b/crates/ruff_linter/resources/test/fixtures/pycodestyle/E402_2.py new file mode 100644 index 00000000000000..f510488cc805f0 --- /dev/null +++ b/crates/ruff_linter/resources/test/fixtures/pycodestyle/E402_2.py @@ -0,0 +1,6 @@ +import os + +os.environ["WORLD_SIZE"] = "1" +os.putenv("CUDA_VISIBLE_DEVICES", "4") + +import torch diff --git a/crates/ruff_linter/src/checkers/ast/mod.rs b/crates/ruff_linter/src/checkers/ast/mod.rs index e858b53899e2fc..df79cf9f087ed5 100644 --- a/crates/ruff_linter/src/checkers/ast/mod.rs +++ b/crates/ruff_linter/src/checkers/ast/mod.rs @@ -374,7 +374,9 @@ where || helpers::is_assignment_to_a_dunder(stmt) || helpers::in_nested_block(self.semantic.current_statements()) || imports::is_matplotlib_activation(stmt, self.semantic()) - || imports::is_sys_path_modification(stmt, self.semantic())) + || imports::is_sys_path_modification(stmt, self.semantic()) + || (self.settings.preview.is_enabled() + && imports::is_os_environ_modification(stmt, self.semantic()))) { self.semantic.flags |= SemanticModelFlags::IMPORT_BOUNDARY; } diff --git a/crates/ruff_linter/src/rules/pycodestyle/mod.rs b/crates/ruff_linter/src/rules/pycodestyle/mod.rs index 5589733bef21db..34b63a26ffbe9a 100644 --- a/crates/ruff_linter/src/rules/pycodestyle/mod.rs +++ b/crates/ruff_linter/src/rules/pycodestyle/mod.rs @@ -38,6 +38,7 @@ mod tests { #[test_case(Rule::ModuleImportNotAtTopOfFile, Path::new("E40.py"))] #[test_case(Rule::ModuleImportNotAtTopOfFile, Path::new("E402_0.py"))] #[test_case(Rule::ModuleImportNotAtTopOfFile, Path::new("E402_1.py"))] + #[test_case(Rule::ModuleImportNotAtTopOfFile, Path::new("E402_2.py"))] #[test_case(Rule::ModuleImportNotAtTopOfFile, Path::new("E402.ipynb"))] #[test_case(Rule::MultipleImportsOnOneLine, Path::new("E40.py"))] #[test_case(Rule::MultipleStatementsOnOneLineColon, Path::new("E70.py"))] @@ -69,6 +70,7 @@ mod tests { #[test_case(Rule::IsLiteral, Path::new("constant_literals.py"))] #[test_case(Rule::TypeComparison, Path::new("E721.py"))] + #[test_case(Rule::ModuleImportNotAtTopOfFile, Path::new("E402_2.py"))] fn preview_rules(rule_code: Rule, path: &Path) -> Result<()> { let snapshot = format!( "preview__{}_{}", diff --git a/crates/ruff_linter/src/rules/pycodestyle/rules/module_import_not_at_top_of_file.rs b/crates/ruff_linter/src/rules/pycodestyle/rules/module_import_not_at_top_of_file.rs index 7e679809fc5f82..1112cd7cc5b509 100644 --- a/crates/ruff_linter/src/rules/pycodestyle/rules/module_import_not_at_top_of_file.rs +++ b/crates/ruff_linter/src/rules/pycodestyle/rules/module_import_not_at_top_of_file.rs @@ -17,6 +17,9 @@ use crate::checkers::ast::Checker; /// `sys.path.insert`, `sys.path.append`, and similar modifications between import /// statements. /// +/// In [preview], this rule also allows `os.environ` modifications between import +/// statements. +/// /// ## Example /// ```python /// "One string" @@ -37,6 +40,7 @@ use crate::checkers::ast::Checker; /// ``` /// /// [PEP 8]: https://peps.python.org/pep-0008/#imports +/// [preview]: https://docs.astral.sh/ruff/preview/ #[violation] pub struct ModuleImportNotAtTopOfFile { source_type: PySourceType, @@ -54,6 +58,7 @@ impl Violation for ModuleImportNotAtTopOfFile { } /// E402 +/// pub(crate) fn module_import_not_at_top_of_file(checker: &mut Checker, stmt: &Stmt) { if checker.semantic().seen_import_boundary() && checker.semantic().at_top_level() { checker.diagnostics.push(Diagnostic::new( diff --git a/crates/ruff_linter/src/rules/pycodestyle/snapshots/ruff_linter__rules__pycodestyle__tests__E402_E402_2.py.snap b/crates/ruff_linter/src/rules/pycodestyle/snapshots/ruff_linter__rules__pycodestyle__tests__E402_E402_2.py.snap new file mode 100644 index 00000000000000..67363450ed4a8a --- /dev/null +++ b/crates/ruff_linter/src/rules/pycodestyle/snapshots/ruff_linter__rules__pycodestyle__tests__E402_E402_2.py.snap @@ -0,0 +1,12 @@ +--- +source: crates/ruff_linter/src/rules/pycodestyle/mod.rs +--- +E402_2.py:6:1: E402 Module level import not at top of file + | +4 | os.putenv("CUDA_VISIBLE_DEVICES", "4") +5 | +6 | import torch + | ^^^^^^^^^^^^ E402 + | + + diff --git a/crates/ruff_linter/src/rules/pycodestyle/snapshots/ruff_linter__rules__pycodestyle__tests__preview__E402_E402_2.py.snap b/crates/ruff_linter/src/rules/pycodestyle/snapshots/ruff_linter__rules__pycodestyle__tests__preview__E402_E402_2.py.snap new file mode 100644 index 00000000000000..6dcc4546f11f9e --- /dev/null +++ b/crates/ruff_linter/src/rules/pycodestyle/snapshots/ruff_linter__rules__pycodestyle__tests__preview__E402_E402_2.py.snap @@ -0,0 +1,4 @@ +--- +source: crates/ruff_linter/src/rules/pycodestyle/mod.rs +--- + diff --git a/crates/ruff_python_semantic/src/analyze/imports.rs b/crates/ruff_python_semantic/src/analyze/imports.rs index a4a9ddc134486d..04aba5e9f63a41 100644 --- a/crates/ruff_python_semantic/src/analyze/imports.rs +++ b/crates/ruff_python_semantic/src/analyze/imports.rs @@ -1,3 +1,4 @@ +use ruff_python_ast::helpers::map_subscript; use ruff_python_ast::{self as ast, Expr, Stmt}; use crate::SemanticModel; @@ -36,6 +37,50 @@ pub fn is_sys_path_modification(stmt: &Stmt, semantic: &SemanticModel) -> bool { }) } +/// Returns `true` if a [`Stmt`] is an `os.environ` modification, as in: +/// ```python +/// import os +/// +/// os.environ["CUDA_VISIBLE_DEVICES"] = "4" +/// ``` +pub fn is_os_environ_modification(stmt: &Stmt, semantic: &SemanticModel) -> bool { + match stmt { + Stmt::Expr(ast::StmtExpr { value, .. }) => match value.as_ref() { + Expr::Call(ast::ExprCall { func, .. }) => semantic + .resolve_call_path(func.as_ref()) + .is_some_and(|call_path| { + matches!( + call_path.as_slice(), + ["os", "putenv" | "unsetenv"] + | [ + "os", + "environ", + "update" | "pop" | "clear" | "setdefault" | "popitem" + ] + ) + }), + _ => false, + }, + Stmt::Delete(ast::StmtDelete { targets, .. }) => targets.iter().any(|target| { + semantic + .resolve_call_path(map_subscript(target)) + .is_some_and(|call_path| matches!(call_path.as_slice(), ["os", "environ"])) + }), + Stmt::Assign(ast::StmtAssign { targets, .. }) => targets.iter().any(|target| { + semantic + .resolve_call_path(map_subscript(target)) + .is_some_and(|call_path| matches!(call_path.as_slice(), ["os", "environ"])) + }), + Stmt::AnnAssign(ast::StmtAnnAssign { target, .. }) => semantic + .resolve_call_path(map_subscript(target)) + .is_some_and(|call_path| matches!(call_path.as_slice(), ["os", "environ"])), + Stmt::AugAssign(ast::StmtAugAssign { target, .. }) => semantic + .resolve_call_path(map_subscript(target)) + .is_some_and(|call_path| matches!(call_path.as_slice(), ["os", "environ"])), + _ => false, + } +} + /// Returns `true` if a [`Stmt`] is a `matplotlib.use` activation, as in: /// ```python /// import matplotlib