forked from astral-sh/ruff
-
Notifications
You must be signed in to change notification settings - Fork 0
/
imports.rs
100 lines (96 loc) · 3.48 KB
/
imports.rs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
use ruff_python_ast::helpers::map_subscript;
use ruff_python_ast::{self as ast, Expr, Stmt};
use crate::SemanticModel;
/// Returns `true` if a [`Stmt`] is a `sys.path` modification, as in:
/// ```python
/// import sys
///
/// sys.path.append("../")
/// ```
pub fn is_sys_path_modification(stmt: &Stmt, semantic: &SemanticModel) -> bool {
let Stmt::Expr(ast::StmtExpr { value, range: _ }) = stmt else {
return false;
};
let Expr::Call(ast::ExprCall { func, .. }) = value.as_ref() else {
return false;
};
semantic
.resolve_call_path(func.as_ref())
.is_some_and(|call_path| {
matches!(
call_path.as_slice(),
[
"sys",
"path",
"append"
| "insert"
| "extend"
| "remove"
| "pop"
| "clear"
| "reverse"
| "sort"
]
)
})
}
/// 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
///
/// matplotlib.use("Agg")
/// ```
pub fn is_matplotlib_activation(stmt: &Stmt, semantic: &SemanticModel) -> bool {
let Stmt::Expr(ast::StmtExpr { value, range: _ }) = stmt else {
return false;
};
let Expr::Call(ast::ExprCall { func, .. }) = value.as_ref() else {
return false;
};
semantic
.resolve_call_path(func.as_ref())
.is_some_and(|call_path| matches!(call_path.as_slice(), ["matplotlib", "use"]))
}