diff --git a/resources/test/fixtures/F821.py b/resources/test/fixtures/F821_0.py similarity index 100% rename from resources/test/fixtures/F821.py rename to resources/test/fixtures/F821_0.py diff --git a/resources/test/fixtures/F821_1.py b/resources/test/fixtures/F821_1.py new file mode 100644 index 0000000000..3a562037e3 --- /dev/null +++ b/resources/test/fixtures/F821_1.py @@ -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) diff --git a/src/ast/helpers.rs b/src/ast/helpers.rs index 1ae3df3267..986b21558b 100644 --- a/src/ast/helpers.rs +++ b/src/ast/helpers.rs @@ -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 { } } +/// 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 { +/// 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, } } diff --git a/src/check_ast.rs b/src/check_ast.rs index 1ff3f0aaea..acf7c19152 100644 --- a/src/check_ast.rs +++ b/src/check_ast.rs @@ -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>, + 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 => { diff --git a/src/linter.rs b/src/linter.rs index 1f8e789cb9..04f45d2c55 100644 --- a/src/linter.rs +++ b/src/linter.rs @@ -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")] diff --git a/src/python/typing.rs b/src/python/typing.rs index 4707db24fe..a0409a3054 100644 --- a/src/python/typing.rs +++ b/src/python/typing.rs @@ -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> = Lazy::new(|| { BTreeSet::from([ "AbstractAsyncContextManager", @@ -14,7 +17,6 @@ static ANNOTATED_SUBSCRIPTS: Lazy> = Lazy::new(|| { "AsyncIterator", "Awaitable", "BinaryIO", - "BsdDbShelf", "ByteString", "Callable", "ChainMap", @@ -72,30 +74,90 @@ static ANNOTATED_SUBSCRIPTS: Lazy> = 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> = + Lazy::new(|| BTreeSet::from(["Dict", "FrozenSet", "List", "Set", "Tuple", "Type"])); + +static PEP_585_BUILTINS: Lazy> = + 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> = - 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 { + 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, + } } diff --git a/src/pyupgrade/plugins/use_pep585_annotation.rs b/src/pyupgrade/plugins/use_pep585_annotation.rs index bcc236af15..3498a437ac 100644 --- a/src/pyupgrade/plugins/use_pep585_annotation.rs +++ b/src/pyupgrade/plugins/use_pep585_annotation.rs @@ -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); } diff --git a/src/pyupgrade/plugins/use_pep604_annotation.rs b/src/pyupgrade/plugins/use_pep604_annotation.rs index 27e01e5f3b..e71c7c5eca 100644 --- a/src/pyupgrade/plugins/use_pep604_annotation.rs +++ b/src/pyupgrade/plugins/use_pep604_annotation.rs @@ -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 { diff --git a/src/snapshots/ruff__linter__tests__F821_F821.py.snap b/src/snapshots/ruff__linter__tests__F821_F821_0.py.snap similarity index 100% rename from src/snapshots/ruff__linter__tests__F821_F821.py.snap rename to src/snapshots/ruff__linter__tests__F821_F821_0.py.snap diff --git a/src/snapshots/ruff__linter__tests__F821_F821_1.py.snap b/src/snapshots/ruff__linter__tests__F821_F821_1.py.snap new file mode 100644 index 0000000000..3c25cf29b6 --- /dev/null +++ b/src/snapshots/ruff__linter__tests__F821_F821_1.py.snap @@ -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: ~ +