Allow os.environ modifications between imports

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

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;
                 }

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__{}_{}",

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(

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
+  |
+
+

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
+---
+

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