mirrors/ruff
1
0
Fork 0
mirror of https://github.com/astral-sh/ruff.git synced 2025-08-04 02:38:25 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

[pycodestyle] Allow os.environ modifications between imports (E402) (#10066)

Browse source 
## 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
This commit is contained in:
Charlie Marsh 2024-02-20 13:24:27 -05:00 • committed by GitHub
parent 7eafba2a4d
commit 4997c681f1
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 77 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

45
crates/ruff_python_semantic/src/analyze/imports.rs
View file

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