Python
/
ruff
mirror of https://github.com/astral-sh/ruff
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Track typing module imports (#533)

This commit is contained in:
Charlie Marsh 2022-11-01 14:01:59 -04:00 committed by GitHub
parent 927d716edd
commit e9a4c8ba13
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
10 changed files with 192 additions and 71 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

0
resources/test/fixtures/F821.py → resources/test/fixtures/F821_0.py vendored
View File

18
resources/test/fixtures/F821_1.py vendored Normal file
View File

@ -0,0 +1,18 @@
"""Test: typing module imports."""
from foo import cast
# OK
x = cast("Model")
from typing import cast
# F821 Undefined name `Model`
x = cast("Model", x)
import typing
# F821 Undefined name `Model`
x = typing.cast("Model", x)

46
src/ast/helpers.rs
View File

@ -1,9 +1,9 @@
use std::collections::BTreeSet;
use once_cell::sync::Lazy;
use regex::Regex;
use rustpython_ast::{Excepthandler, ExcepthandlerKind, Expr, ExprKind, Location, StmtKind};
use crate::python::typing;
fn compose_call_path_inner<'a>(expr: &'a Expr, parts: &mut Vec<&'a str>) {
match &expr.node {
ExprKind::Call { func, .. } => {
@ -30,6 +30,7 @@ pub fn compose_call_path(expr: &Expr) -> Option<String> {
}
}
/// Return `true` if the `Expr` is a name or attribute reference to `${target}`.
pub fn match_name_or_attr(expr: &Expr, target: &str) -> bool {
match &expr.node {
ExprKind::Attribute { attr, .. } => target == attr,
@ -38,32 +39,27 @@ pub fn match_name_or_attr(expr: &Expr, target: &str) -> bool {
}
}
pub enum SubscriptKind {
AnnotatedSubscript,
PEP593AnnotatedSubscript,
}
pub fn match_annotated_subscript(expr: &Expr) -> Option<SubscriptKind> {
/// Return `true` if the `Expr` is a reference to `${module}.${target}`.
///
/// Useful for, e.g., ensuring that a `Union` reference represents `typing.Union`.
pub fn match_name_or_attr_from_module(
expr: &Expr,
target: &str,
module: &str,
imports: Option<&BTreeSet<&str>>,
) -> bool {
match &expr.node {
ExprKind::Attribute { attr, .. } => {
if typing::is_annotated_subscript(attr) {
Some(SubscriptKind::AnnotatedSubscript)
} else if typing::is_pep593_annotated_subscript(attr) {
Some(SubscriptKind::PEP593AnnotatedSubscript)
} else {
None
}
}
ExprKind::Attribute { value, attr, .. } => match &value.node {
ExprKind::Name { id, .. } => id == module && target == attr,
_ => false,
},
ExprKind::Name { id, .. } => {
if typing::is_annotated_subscript(id) {
Some(SubscriptKind::AnnotatedSubscript)
} else if typing::is_pep593_annotated_subscript(id) {
Some(SubscriptKind::PEP593AnnotatedSubscript)
} else {
None
}
target == id
&& imports
.map(|imports| imports.contains(&id.as_str()))
.unwrap_or_default()
}
_ => None,
_ => false,
}
}

60
src/check_ast.rs
View File

@ -11,7 +11,7 @@ use rustpython_parser::ast::{
};
use rustpython_parser::parser;
use crate::ast::helpers::{extract_handler_names, match_name_or_attr, SubscriptKind};
use crate::ast::helpers::{extract_handler_names, match_name_or_attr_from_module};
use crate::ast::operations::extract_all_names;
use crate::ast::relocate::relocate_expr;
use crate::ast::types::{
@ -25,6 +25,8 @@ use crate::checks::{Check, CheckCode, CheckKind};
use crate::docstrings::definition::{Definition, DefinitionKind, Documentable};
use crate::python::builtins::{BUILTINS, MAGIC_GLOBALS};
use crate::python::future::ALL_FEATURE_NAMES;
use crate::python::typing;
use crate::python::typing::SubscriptKind;
use crate::settings::types::PythonVersion;
use crate::settings::Settings;
use crate::source_code_locator::SourceCodeLocator;
@ -34,7 +36,8 @@ use crate::{
pycodestyle, pydocstyle, pyflakes, pyupgrade,
};
pub const GLOBAL_SCOPE_INDEX: usize = 0;
const GLOBAL_SCOPE_INDEX: usize = 0;
const TRACK_FROM_IMPORTS: [&str; 1] = ["typing"];
pub struct Checker<'a> {
// Input data.
@ -70,6 +73,7 @@ pub struct Checker<'a> {
futures_allowed: bool,
annotations_future_enabled: bool,
except_handlers: Vec<Vec<String>>,
from_imports: BTreeMap<&'a str, BTreeSet<&'a str>>,
}
impl<'a> Checker<'a> {
@ -101,13 +105,14 @@ impl<'a> Checker<'a> {
modifier: Modifier::Module,
visibility: module_visibility(path),
},
in_f_string: None,
in_f_string: Default::default(),
in_annotation: Default::default(),
in_literal: Default::default(),
seen_import_boundary: Default::default(),
futures_allowed: true,
annotations_future_enabled: Default::default(),
except_handlers: Default::default(),
from_imports: Default::default(),
}
}
@ -115,6 +120,11 @@ impl<'a> Checker<'a> {
pub fn patch(&self) -> bool {
self.autofix.patch()
}
/// Return `true` if the `Expr` is a reference to `typing.${target}`.
pub fn match_typing_module(&self, expr: &Expr, target: &str) -> bool {
match_name_or_attr_from_module(expr, target, "typing", self.from_imports.get("typing"))
}
}
impl<'a, 'b> Visitor<'b> for Checker<'a>
@ -492,6 +502,24 @@ where
module,
level,
} => {
// Track `import from` statements, to ensure that we can correctly attribute
// references like `from typing import Union`.
if level.map(|level| level == 0).unwrap_or(true) {
if let Some(module) = module {
if TRACK_FROM_IMPORTS.contains(&module.as_str()) {
self.from_imports
.entry(module)
.or_insert_with(BTreeSet::new)
.extend(
names
.iter()
.filter(|alias| alias.node.asname.is_none())
.map(|alias| alias.node.name.as_str()),
)
}
}
}
if self.settings.enabled.contains(&CheckCode::E402) {
if self.seen_import_boundary && stmt.location.column() == 0 {
self.checks.push(Check::new(
@ -861,7 +889,7 @@ where
pyupgrade::plugins::use_pep604_annotation(self, expr, value, slice);
}
if match_name_or_attr(value, "Literal") {
if self.match_typing_module(value, "Literal") {
self.in_literal = true;
}
}
@ -886,6 +914,7 @@ where
// Ex) List[...]
if self.settings.enabled.contains(&CheckCode::U006)
&& self.settings.target_version >= PythonVersion::Py39
&& typing::is_pep585_builtin(expr, self.from_imports.get("typing"))
{
pyupgrade::plugins::use_pep585_annotation(self, expr, id);
}
@ -908,16 +937,13 @@ where
}
ExprContext::Del => self.handle_node_delete(expr),
},
ExprKind::Attribute { value, attr, .. } => {
ExprKind::Attribute { attr, .. } => {
// Ex) typing.List[...]
if self.settings.enabled.contains(&CheckCode::U006)
&& self.settings.target_version >= PythonVersion::Py39
&& typing::is_pep585_builtin(expr, self.from_imports.get("typing"))
{
if let ExprKind::Name { id, .. } = &value.node {
if id == "typing" {
pyupgrade::plugins::use_pep585_annotation(self, expr, attr);
}
}
pyupgrade::plugins::use_pep585_annotation(self, expr, attr);
}
}
ExprKind::Call {
@ -1276,12 +1302,12 @@ where
args,
keywords,
} => {
if match_name_or_attr(func, "ForwardRef") {
if self.match_typing_module(func, "ForwardRef") {
self.visit_expr(func);
for expr in args {
self.visit_annotation(expr);
}
} else if match_name_or_attr(func, "cast") {
} else if self.match_typing_module(func, "cast") {
self.visit_expr(func);
if !args.is_empty() {
self.visit_annotation(&args[0]);
@ -1289,12 +1315,12 @@ where
for expr in args.iter().skip(1) {
self.visit_expr(expr);
}
} else if match_name_or_attr(func, "NewType") {
} else if self.match_typing_module(func, "NewType") {
self.visit_expr(func);
for expr in args.iter().skip(1) {
self.visit_annotation(expr);
}
} else if match_name_or_attr(func, "TypeVar") {
} else if self.match_typing_module(func, "TypeVar") {
self.visit_expr(func);
for expr in args.iter().skip(1) {
self.visit_annotation(expr);
@ -1311,7 +1337,7 @@ where
}
}
}
} else if match_name_or_attr(func, "NamedTuple") {
} else if self.match_typing_module(func, "NamedTuple") {
self.visit_expr(func);
// Ex) NamedTuple("a", [("a", int)])
@ -1343,7 +1369,7 @@ where
let KeywordData { value, .. } = &keyword.node;
self.visit_annotation(value);
}
} else if match_name_or_attr(func, "TypedDict") {
} else if self.match_typing_module(func, "TypedDict") {
self.visit_expr(func);
// Ex) TypedDict("a", {"a": int})
@ -1370,7 +1396,7 @@ where
}
}
ExprKind::Subscript { value, slice, ctx } => {
match helpers::match_annotated_subscript(value) {
match typing::match_annotated_subscript(value, self.from_imports.get("typing")) {
Some(subscript) => match subscript {
// Ex) Optional[int]
SubscriptKind::AnnotatedSubscript => {

3
src/linter.rs
View File

@ -377,7 +377,8 @@ mod tests {
#[test_case(CheckCode::F706, Path::new("F706.py"); "F706")]
#[test_case(CheckCode::F707, Path::new("F707.py"); "F707")]
#[test_case(CheckCode::F722, Path::new("F722.py"); "F722")]
#[test_case(CheckCode::F821, Path::new("F821.py"); "F821")]
#[test_case(CheckCode::F821, Path::new("F821_0.py"); "F821_0")]
#[test_case(CheckCode::F821, Path::new("F821_1.py"); "F821_1")]
#[test_case(CheckCode::F822, Path::new("F822.py"); "F822")]
#[test_case(CheckCode::F823, Path::new("F823.py"); "F823")]
#[test_case(CheckCode::F831, Path::new("F831.py"); "F831")]

82
src/python/typing.rs
View File

@ -1,7 +1,10 @@
use std::collections::BTreeSet;
use once_cell::sync::Lazy;
use rustpython_ast::{Expr, ExprKind};
// TODO(charlie): Some of these are actually from `collections`.
// Review: https://peps.python.org/pep-0585/.
static ANNOTATED_SUBSCRIPTS: Lazy<BTreeSet<&'static str>> = Lazy::new(|| {
BTreeSet::from([
"AbstractAsyncContextManager",
@ -14,7 +17,6 @@ static ANNOTATED_SUBSCRIPTS: Lazy<BTreeSet<&'static str>> = Lazy::new(|| {
"AsyncIterator",
"Awaitable",
"BinaryIO",
"BsdDbShelf",
"ByteString",
"Callable",
"ChainMap",
@ -72,30 +74,90 @@ static ANNOTATED_SUBSCRIPTS: Lazy<BTreeSet<&'static str>> = Lazy::new(|| {
"WeakMethod",
"WeakSet",
"WeakValueDictionary",
"cached_property",
"defaultdict",
"deque",
"dict",
"frozenset",
"list",
"partialmethod",
"set",
"tuple",
"type",
])
});
pub fn is_annotated_subscript(name: &str) -> bool {
static PEP_585_BUILTINS_ELIGIBLE: Lazy<BTreeSet<&'static str>> =
Lazy::new(|| BTreeSet::from(["Dict", "FrozenSet", "List", "Set", "Tuple", "Type"]));
static PEP_585_BUILTINS: Lazy<BTreeSet<&'static str>> =
Lazy::new(|| BTreeSet::from(["dict", "frozenset", "list", "set", "tuple", "type"]));
fn is_annotated_subscript(name: &str) -> bool {
ANNOTATED_SUBSCRIPTS.contains(name)
}
pub fn is_pep593_annotated_subscript(name: &str) -> bool {
fn is_pep593_annotated_subscript(name: &str) -> bool {
name == "Annotated"
}
static PEP_585_BUILTINS: Lazy<BTreeSet<&'static str>> =
Lazy::new(|| BTreeSet::from(["Dict", "FrozenSet", "List", "Set", "Tuple", "Type"]));
pub fn is_pep585_builtin(name: &str) -> bool {
PEP_585_BUILTINS.contains(name)
pub enum SubscriptKind {
AnnotatedSubscript,
PEP593AnnotatedSubscript,
}
pub fn match_annotated_subscript(
expr: &Expr,
imports: Option<&BTreeSet<&str>>,
) -> Option<SubscriptKind> {
match &expr.node {
ExprKind::Attribute { attr, value, .. } => {
if let ExprKind::Name { id, .. } = &value.node {
if id == "typing" {
if is_pep593_annotated_subscript(attr) {
return Some(SubscriptKind::PEP593AnnotatedSubscript);
} else if is_annotated_subscript(attr) {
return Some(SubscriptKind::AnnotatedSubscript);
}
}
}
}
ExprKind::Name { id, .. } => {
// Built-ins (no import necessary).
if PEP_585_BUILTINS.contains(&id.as_str()) {
return Some(SubscriptKind::AnnotatedSubscript);
}
if imports
.map(|import| import.contains(&id.as_str()))
.unwrap_or_default()
{
if is_pep593_annotated_subscript(id) {
return Some(SubscriptKind::PEP593AnnotatedSubscript);
} else if is_annotated_subscript(id) {
return Some(SubscriptKind::AnnotatedSubscript);
}
}
}
_ => {}
}
None
}
/// Returns `true` if `Expr` represents a reference to a typing object with a PEP585 built-in.
pub fn is_pep585_builtin(expr: &Expr, imports: Option<&BTreeSet<&str>>) -> bool {
match &expr.node {
ExprKind::Attribute { attr, value, .. } => {
if let ExprKind::Name { id, .. } = &value.node {
id == "typing" && PEP_585_BUILTINS_ELIGIBLE.contains(&attr.as_str())
} else {
false
}
}
ExprKind::Name { id, .. } => {
imports
.map(|import| import.contains(&id.as_str()))
.unwrap_or_default()
&& PEP_585_BUILTINS_ELIGIBLE.contains(&id.as_str())
}
_ => false,
}
}

26
src/pyupgrade/plugins/use_pep585_annotation.rs
View File

@ -4,22 +4,18 @@ use crate::ast::types::Range;
use crate::autofix::Fix;
use crate::check_ast::Checker;
use crate::checks::{Check, CheckKind};
use crate::python::typing;
pub fn use_pep585_annotation(checker: &mut Checker, expr: &Expr, id: &str) {
// TODO(charlie): Verify that the builtin is imported from the `typing` module.
if typing::is_pep585_builtin(id) {
let mut check = Check::new(
CheckKind::UsePEP585Annotation(id.to_string()),
Range::from_located(expr),
);
if checker.patch() {
check.amend(Fix::replacement(
id.to_lowercase(),
expr.location,
expr.end_location.unwrap(),
));
}
checker.add_check(check);
let mut check = Check::new(
CheckKind::UsePEP585Annotation(id.to_string()),
Range::from_located(expr),
);
if checker.patch() {
check.amend(Fix::replacement(
id.to_lowercase(),
expr.location,
expr.end_location.unwrap(),
));
}
checker.add_check(check);
}

5
src/pyupgrade/plugins/use_pep604_annotation.rs
View File

@ -1,6 +1,5 @@
use rustpython_ast::{Constant, Expr, ExprKind, Operator};
use crate::ast::helpers::match_name_or_attr;
use crate::ast::types::Range;
use crate::autofix::Fix;
use crate::check_ast::Checker;
@ -43,7 +42,7 @@ fn union(elts: &[Expr]) -> Expr {
}
pub fn use_pep604_annotation(checker: &mut Checker, expr: &Expr, value: &Expr, slice: &Expr) {
if match_name_or_attr(value, "Optional") {
if checker.match_typing_module(value, "Optional") {
let mut check = Check::new(CheckKind::UsePEP604Annotation, Range::from_located(expr));
if checker.patch() {
let mut generator = SourceGenerator::new();
@ -58,7 +57,7 @@ pub fn use_pep604_annotation(checker: &mut Checker, expr: &Expr, value: &Expr, s
}
}
checker.add_check(check);
} else if match_name_or_attr(value, "Union") {
} else if checker.match_typing_module(value, "Union") {
let mut check = Check::new(CheckKind::UsePEP604Annotation, Range::from_located(expr));
if checker.patch() {
match &slice.node {

0
src/snapshots/ruff__linter__tests__F821_F821.py.snap → src/snapshots/ruff__linter__tests__F821_F821_0.py.snap
View File

23
src/snapshots/ruff__linter__tests__F821_F821_1.py.snap Normal file
View File

@ -0,0 +1,23 @@
---
source: src/linter.rs
expression: checks
---
- kind:
UndefinedName: Model
location:
row: 11
column: 9
end_location:
row: 11
column: 16
fix: ~
- kind:
UndefinedName: Model
location:
row: 18
column: 16
end_location:
row: 18
column: 23
fix: ~