From aa59a370c6f7843f8bd112b2f3c46a8fc599b621 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Tue, 20 Feb 2024 13:24:27 -0500 Subject: [PATCH] [`pycodestyle`] Allow `os.environ` modifications between imports (`E402`) (#10066) ## Summary Allows, e.g.: ```python import os os.environ["WORLD_SIZE"] = "1" os.putenv("CUDA_VISIBLE_DEVICES", "4") import torch ``` For now, this is only allowed in preview. Closes https://github.com/astral-sh/ruff/issues/10059 --- .../test/fixtures/pycodestyle/E402_2.py | 7 +++ 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 | 4 ++ ...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 0000000000000..3579c45babe2e --- /dev/null +++ b/crates/ruff_linter/resources/test/fixtures/pycodestyle/E402_2.py @@ -0,0 +1,7 @@ +import os + +os.environ["WORLD_SIZE"] = "1" +os.putenv("CUDA_VISIBLE_DEVICES", "4") +del os.environ["WORLD_SIZE"] + +import torch diff --git a/crates/ruff_linter/src/checkers/ast/mod.rs b/crates/ruff_linter/src/checkers/ast/mod.rs index e858b53899e2f..df79cf9f087ed 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 5589733bef21d..34b63a26ffbe9 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 7e679809fc5f8..9d2a9bbd61bef 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, 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 0000000000000..59c85e39a83e4 --- /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:7:1: E402 Module level import not at top of file + | +5 | del os.environ["WORLD_SIZE"] +6 | +7 | 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 0000000000000..6dcc4546f11f9 --- /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 a4a9ddc134486..04aba5e9f63a4 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