Skip to content

Commit

Permalink
Allow os.environ modifications between imports
Browse files Browse the repository at this point in the history
  • Loading branch information
charliermarsh committed Feb 20, 2024
1 parent ee4efdb commit f6e49cf
Show file tree
Hide file tree
Showing 7 changed files with 77 additions and 1 deletion.
Original file line number Diff line number Diff line change
@@ -0,0 +1,6 @@
import os

os.environ["WORLD_SIZE"] = "1"
os.putenv("CUDA_VISIBLE_DEVICES", "4")

import torch
4 changes: 3 additions & 1 deletion crates/ruff_linter/src/checkers/ast/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
Expand Down
2 changes: 2 additions & 0 deletions crates/ruff_linter/src/rules/pycodestyle/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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"))]
Expand Down Expand Up @@ -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__{}_{}",
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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,
Expand All @@ -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(
Expand Down
Original file line number Diff line number Diff line change
@@ -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
|


Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
---
source: crates/ruff_linter/src/rules/pycodestyle/mod.rs
---

45 changes: 45 additions & 0 deletions crates/ruff_python_semantic/src/analyze/imports.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
use ruff_python_ast::helpers::map_subscript;
use ruff_python_ast::{self as ast, Expr, Stmt};

use crate::SemanticModel;
Expand Down Expand Up @@ -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
Expand Down

0 comments on commit f6e49cf

Please sign in to comment.