From 1a65e544c5896634a00c05bc463d0e60ec737001 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Tue, 12 Dec 2023 22:12:38 -0500 Subject: [PATCH] Allow `flake8-type-checking` rules to automatically quote runtime-evaluated references (#6001) ## Summary This allows us to fix usages like: ```python from pandas import DataFrame def baz() -> DataFrame: ... ``` By quoting the `DataFrame` in `-> DataFrame`. Without quotes, moving `from pandas import DataFrame` into an `if TYPE_CHECKING:` block will fail at runtime, since Python tries to evaluate the annotation to add it to the function's `__annotations__`. Unfortunately, this does require us to split our "annotation kind" flags into three categories, rather than two: - `typing-only`: The annotation is only evaluated at type-checking-time. - `runtime-evaluated`: Python will evaluate the annotation at runtime (like above) -- but we're willing to quote it. - `runtime-required`: Python will evaluate the annotation at runtime (like above), and some library (like Pydantic) needs it to be available at runtime, so we _can't_ quote it. This functionality is gated behind a setting (`flake8-type-checking.quote-annotations`). Closes https://github.com/astral-sh/ruff/issues/5559. --- .../fixtures/flake8_type_checking/quote.py | 67 +++++ .../checkers/ast/analyze/deferred_scopes.rs | 1 + .../src/checkers/ast/annotation.rs | 66 +++++ crates/ruff_linter/src/checkers/ast/mod.rs | 68 ++--- .../src/rules/flake8_type_checking/helpers.rs | 120 +++++++- .../src/rules/flake8_type_checking/imports.rs | 22 ++ .../src/rules/flake8_type_checking/mod.rs | 27 +- .../runtime_import_in_type_checking_block.rs | 258 +++++++++++++----- .../rules/typing_only_runtime_import.rs | 86 +++--- .../rules/flake8_type_checking/settings.rs | 12 +- ...mport-in-type-checking-block_quote.py.snap | 22 ++ ...ping-only-third-party-import_quote.py.snap | 199 ++++++++++++++ ...mport-in-type-checking-block_quote.py.snap | 29 ++ ...ping-only-third-party-import_quote.py.snap | 4 + crates/ruff_python_semantic/src/model.rs | 142 +++++++--- crates/ruff_python_semantic/src/reference.rs | 54 +++- crates/ruff_workspace/src/options.rs | 56 +++- ruff.schema.json | 7 + 18 files changed, 1033 insertions(+), 207 deletions(-) create mode 100644 crates/ruff_linter/resources/test/fixtures/flake8_type_checking/quote.py create mode 100644 crates/ruff_linter/src/checkers/ast/annotation.rs create mode 100644 crates/ruff_linter/src/rules/flake8_type_checking/imports.rs create mode 100644 crates/ruff_linter/src/rules/flake8_type_checking/snapshots/ruff_linter__rules__flake8_type_checking__tests__quote_runtime-import-in-type-checking-block_quote.py.snap create mode 100644 crates/ruff_linter/src/rules/flake8_type_checking/snapshots/ruff_linter__rules__flake8_type_checking__tests__quote_typing-only-third-party-import_quote.py.snap create mode 100644 crates/ruff_linter/src/rules/flake8_type_checking/snapshots/ruff_linter__rules__flake8_type_checking__tests__runtime-import-in-type-checking-block_quote.py.snap create mode 100644 crates/ruff_linter/src/rules/flake8_type_checking/snapshots/ruff_linter__rules__flake8_type_checking__tests__typing-only-third-party-import_quote.py.snap diff --git a/crates/ruff_linter/resources/test/fixtures/flake8_type_checking/quote.py b/crates/ruff_linter/resources/test/fixtures/flake8_type_checking/quote.py new file mode 100644 index 0000000000000..67332e3010f24 --- /dev/null +++ b/crates/ruff_linter/resources/test/fixtures/flake8_type_checking/quote.py @@ -0,0 +1,67 @@ +def f(): + from pandas import DataFrame + + def baz() -> DataFrame: + ... + + +def f(): + from pandas import DataFrame + + def baz() -> DataFrame[int]: + ... + + +def f(): + from pandas import DataFrame + + def baz() -> DataFrame["int"]: + ... + + +def f(): + import pandas as pd + + def baz() -> pd.DataFrame: + ... + + +def f(): + import pandas as pd + + def baz() -> pd.DataFrame.Extra: + ... + + +def f(): + import pandas as pd + + def baz() -> pd.DataFrame | int: + ... + + + +def f(): + from pandas import DataFrame + + def baz() -> DataFrame(): + ... + + +def f(): + from typing import Literal + + from pandas import DataFrame + + def baz() -> DataFrame[Literal["int"]]: + ... + + +def f(): + from typing import TYPE_CHECKING + + if TYPE_CHECKING: + from pandas import DataFrame + + def func(value: DataFrame): + ... diff --git a/crates/ruff_linter/src/checkers/ast/analyze/deferred_scopes.rs b/crates/ruff_linter/src/checkers/ast/analyze/deferred_scopes.rs index 27c9a6b7c79e2..66d0ad25273e9 100644 --- a/crates/ruff_linter/src/checkers/ast/analyze/deferred_scopes.rs +++ b/crates/ruff_linter/src/checkers/ast/analyze/deferred_scopes.rs @@ -59,6 +59,7 @@ pub(crate) fn deferred_scopes(checker: &mut Checker) { flake8_type_checking::helpers::is_valid_runtime_import( binding, &checker.semantic, + &checker.settings.flake8_type_checking, ) }) .collect() diff --git a/crates/ruff_linter/src/checkers/ast/annotation.rs b/crates/ruff_linter/src/checkers/ast/annotation.rs new file mode 100644 index 0000000000000..aca5fc6c629ff --- /dev/null +++ b/crates/ruff_linter/src/checkers/ast/annotation.rs @@ -0,0 +1,66 @@ +use ruff_python_semantic::{ScopeKind, SemanticModel}; + +use crate::rules::flake8_type_checking; +use crate::settings::LinterSettings; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) enum AnnotationContext { + /// Python will evaluate the annotation at runtime, but it's not _required_ and, as such, could + /// be quoted to convert it into a typing-only annotation. + /// + /// For example: + /// ```python + /// from pandas import DataFrame + /// + /// def foo() -> DataFrame: + /// ... + /// ``` + /// + /// Above, Python will evaluate `DataFrame` at runtime in order to add it to `__annotations__`. + RuntimeEvaluated, + /// Python will evaluate the annotation at runtime, and it's required to be available at + /// runtime, as a library (like Pydantic) needs access to it. + RuntimeRequired, + /// The annotation is only evaluated at type-checking time. + TypingOnly, +} + +impl AnnotationContext { + pub(super) fn from_model(semantic: &SemanticModel, settings: &LinterSettings) -> Self { + // If the annotation is in a class scope (e.g., an annotated assignment for a + // class field), and that class is marked as annotation as runtime-required. + if semantic + .current_scope() + .kind + .as_class() + .is_some_and(|class_def| { + flake8_type_checking::helpers::runtime_required_class( + class_def, + &settings.flake8_type_checking.runtime_required_base_classes, + &settings.flake8_type_checking.runtime_required_decorators, + semantic, + ) + }) + { + return Self::RuntimeRequired; + } + + // If `__future__` annotations are enabled, then annotations are never evaluated + // at runtime, so we can treat them as typing-only. + if semantic.future_annotations() { + return Self::TypingOnly; + } + + // Otherwise, if we're in a class or module scope, then the annotation needs to + // be available at runtime. + // See: https://docs.python.org/3/reference/simple_stmts.html#annotated-assignment-statements + if matches!( + semantic.current_scope().kind, + ScopeKind::Class(_) | ScopeKind::Module + ) { + return Self::RuntimeEvaluated; + } + + Self::TypingOnly + } +} diff --git a/crates/ruff_linter/src/checkers/ast/mod.rs b/crates/ruff_linter/src/checkers/ast/mod.rs index 5d525550f5375..1c510f1cbb8d9 100644 --- a/crates/ruff_linter/src/checkers/ast/mod.rs +++ b/crates/ruff_linter/src/checkers/ast/mod.rs @@ -58,6 +58,7 @@ use ruff_python_semantic::{ use ruff_python_stdlib::builtins::{IPYTHON_BUILTINS, MAGIC_GLOBALS, PYTHON_BUILTINS}; use ruff_source_file::Locator; +use crate::checkers::ast::annotation::AnnotationContext; use crate::checkers::ast::deferred::Deferred; use crate::docstrings::extraction::ExtractionTarget; use crate::importer::Importer; @@ -68,6 +69,7 @@ use crate::settings::{flags, LinterSettings}; use crate::{docstrings, noqa}; mod analyze; +mod annotation; mod deferred; pub(crate) struct Checker<'a> { @@ -515,8 +517,10 @@ where .chain(¶meters.kwonlyargs) { if let Some(expr) = ¶meter_with_default.parameter.annotation { - if runtime_annotation || singledispatch { - self.visit_runtime_annotation(expr); + if singledispatch { + self.visit_runtime_required_annotation(expr); + } else if runtime_annotation { + self.visit_runtime_evaluated_annotation(expr); } else { self.visit_annotation(expr); }; @@ -529,7 +533,7 @@ where if let Some(arg) = ¶meters.vararg { if let Some(expr) = &arg.annotation { if runtime_annotation { - self.visit_runtime_annotation(expr); + self.visit_runtime_evaluated_annotation(expr); } else { self.visit_annotation(expr); }; @@ -538,7 +542,7 @@ where if let Some(arg) = ¶meters.kwarg { if let Some(expr) = &arg.annotation { if runtime_annotation { - self.visit_runtime_annotation(expr); + self.visit_runtime_evaluated_annotation(expr); } else { self.visit_annotation(expr); }; @@ -546,7 +550,7 @@ where } for expr in returns { if runtime_annotation { - self.visit_runtime_annotation(expr); + self.visit_runtime_evaluated_annotation(expr); } else { self.visit_annotation(expr); }; @@ -677,40 +681,16 @@ where value, .. }) => { - // If we're in a class or module scope, then the annotation needs to be - // available at runtime. - // See: https://docs.python.org/3/reference/simple_stmts.html#annotated-assignment-statements - let runtime_annotation = if self.semantic.future_annotations() { - self.semantic - .current_scope() - .kind - .as_class() - .is_some_and(|class_def| { - flake8_type_checking::helpers::runtime_evaluated_class( - class_def, - &self - .settings - .flake8_type_checking - .runtime_evaluated_base_classes, - &self - .settings - .flake8_type_checking - .runtime_evaluated_decorators, - &self.semantic, - ) - }) - } else { - matches!( - self.semantic.current_scope().kind, - ScopeKind::Class(_) | ScopeKind::Module - ) - }; - - if runtime_annotation { - self.visit_runtime_annotation(annotation); - } else { - self.visit_annotation(annotation); + match AnnotationContext::from_model(&self.semantic, self.settings) { + AnnotationContext::RuntimeRequired => { + self.visit_runtime_required_annotation(annotation); + } + AnnotationContext::RuntimeEvaluated => { + self.visit_runtime_evaluated_annotation(annotation); + } + AnnotationContext::TypingOnly => self.visit_annotation(annotation), } + if let Some(expr) = value { if self.semantic.match_typing_expr(annotation, "TypeAlias") { self.visit_type_definition(expr); @@ -1527,10 +1507,18 @@ impl<'a> Checker<'a> { self.semantic.flags = snapshot; } + /// Visit an [`Expr`], and treat it as a runtime-evaluated type annotation. + fn visit_runtime_evaluated_annotation(&mut self, expr: &'a Expr) { + let snapshot = self.semantic.flags; + self.semantic.flags |= SemanticModelFlags::RUNTIME_EVALUATED_ANNOTATION; + self.visit_type_definition(expr); + self.semantic.flags = snapshot; + } + /// Visit an [`Expr`], and treat it as a runtime-required type annotation. - fn visit_runtime_annotation(&mut self, expr: &'a Expr) { + fn visit_runtime_required_annotation(&mut self, expr: &'a Expr) { let snapshot = self.semantic.flags; - self.semantic.flags |= SemanticModelFlags::RUNTIME_ANNOTATION; + self.semantic.flags |= SemanticModelFlags::RUNTIME_REQUIRED_ANNOTATION; self.visit_type_definition(expr); self.semantic.flags = snapshot; } diff --git a/crates/ruff_linter/src/rules/flake8_type_checking/helpers.rs b/crates/ruff_linter/src/rules/flake8_type_checking/helpers.rs index 0a51e151f4703..1fc4ade6fda9a 100644 --- a/crates/ruff_linter/src/rules/flake8_type_checking/helpers.rs +++ b/crates/ruff_linter/src/rules/flake8_type_checking/helpers.rs @@ -1,10 +1,35 @@ +use anyhow::Result; +use rustc_hash::FxHashSet; + +use ruff_diagnostics::Edit; use ruff_python_ast::call_path::from_qualified_name; use ruff_python_ast::helpers::{map_callable, map_subscript}; use ruff_python_ast::{self as ast, Expr}; -use ruff_python_semantic::{Binding, BindingId, BindingKind, SemanticModel}; -use rustc_hash::FxHashSet; +use ruff_python_codegen::Stylist; +use ruff_python_semantic::{ + Binding, BindingId, BindingKind, NodeId, ResolvedReference, SemanticModel, +}; +use ruff_source_file::Locator; +use ruff_text_size::Ranged; + +use crate::rules::flake8_type_checking::settings::Settings; -pub(crate) fn is_valid_runtime_import(binding: &Binding, semantic: &SemanticModel) -> bool { +/// Returns `true` if the [`ResolvedReference`] is in a typing-only context _or_ a runtime-evaluated +/// context (with quoting enabled). +pub(crate) fn is_typing_reference(reference: &ResolvedReference, settings: &Settings) -> bool { + reference.in_type_checking_block() + || reference.in_typing_only_annotation() + || reference.in_complex_string_type_definition() + || reference.in_simple_string_type_definition() + || (settings.quote_annotations && reference.in_runtime_evaluated_annotation()) +} + +/// Returns `true` if the [`Binding`] represents a runtime-required import. +pub(crate) fn is_valid_runtime_import( + binding: &Binding, + semantic: &SemanticModel, + settings: &Settings, +) -> bool { if matches!( binding.kind, BindingKind::Import(..) | BindingKind::FromImport(..) | BindingKind::SubmoduleImport(..) @@ -12,28 +37,29 @@ pub(crate) fn is_valid_runtime_import(binding: &Binding, semantic: &SemanticMode binding.context.is_runtime() && binding .references() - .any(|reference_id| semantic.reference(reference_id).context().is_runtime()) + .map(|reference_id| semantic.reference(reference_id)) + .any(|reference| !is_typing_reference(reference, settings)) } else { false } } -pub(crate) fn runtime_evaluated_class( +pub(crate) fn runtime_required_class( class_def: &ast::StmtClassDef, base_classes: &[String], decorators: &[String], semantic: &SemanticModel, ) -> bool { - if runtime_evaluated_base_class(class_def, base_classes, semantic) { + if runtime_required_base_class(class_def, base_classes, semantic) { return true; } - if runtime_evaluated_decorators(class_def, decorators, semantic) { + if runtime_required_decorators(class_def, decorators, semantic) { return true; } false } -fn runtime_evaluated_base_class( +fn runtime_required_base_class( class_def: &ast::StmtClassDef, base_classes: &[String], semantic: &SemanticModel, @@ -45,7 +71,7 @@ fn runtime_evaluated_base_class( seen: &mut FxHashSet, ) -> bool { class_def.bases().iter().any(|expr| { - // If the base class is itself runtime-evaluated, then this is too. + // If the base class is itself runtime-required, then this is too. // Ex) `class Foo(BaseModel): ...` if semantic .resolve_call_path(map_subscript(expr)) @@ -58,7 +84,7 @@ fn runtime_evaluated_base_class( return true; } - // If the base class extends a runtime-evaluated class, then this does too. + // If the base class extends a runtime-required class, then this does too. // Ex) `class Bar(BaseModel): ...; class Foo(Bar): ...` if let Some(id) = semantic.lookup_attribute(map_subscript(expr)) { if seen.insert(id) { @@ -86,7 +112,7 @@ fn runtime_evaluated_base_class( inner(class_def, base_classes, semantic, &mut FxHashSet::default()) } -fn runtime_evaluated_decorators( +fn runtime_required_decorators( class_def: &ast::StmtClassDef, decorators: &[String], semantic: &SemanticModel, @@ -174,3 +200,75 @@ pub(crate) fn is_singledispatch_implementation( is_singledispatch_interface(function_def, semantic) }) } + +/// Wrap a type annotation in quotes. +/// +/// This requires more than just wrapping the reference itself in quotes. For example: +/// - When quoting `Series` in `Series[pd.Timestamp]`, we want `"Series[pd.Timestamp]"`. +/// - When quoting `kubernetes` in `kubernetes.SecurityContext`, we want `"kubernetes.SecurityContext"`. +/// - When quoting `Series` in `Series["pd.Timestamp"]`, we want `"Series[pd.Timestamp]"`. (This is currently unsupported.) +/// - When quoting `Series` in `Series[Literal["pd.Timestamp"]]`, we want `"Series[Literal['pd.Timestamp']]"`. (This is currently unsupported.) +/// +/// In general, when expanding a component of a call chain, we want to quote the entire call chain. +pub(crate) fn quote_annotation( + node_id: NodeId, + semantic: &SemanticModel, + locator: &Locator, + stylist: &Stylist, +) -> Result { + let expr = semantic.expression(node_id).expect("Expression not found"); + if let Some(parent_id) = semantic.parent_expression_id(node_id) { + match semantic.expression(parent_id) { + Some(Expr::Subscript(parent)) => { + if expr == parent.value.as_ref() { + // If we're quoting the value of a subscript, we need to quote the entire + // expression. For example, when quoting `DataFrame` in `DataFrame[int]`, we + // should generate `"DataFrame[int]"`. + return quote_annotation(parent_id, semantic, locator, stylist); + } + } + Some(Expr::Attribute(parent)) => { + if expr == parent.value.as_ref() { + // If we're quoting the value of an attribute, we need to quote the entire + // expression. For example, when quoting `DataFrame` in `pd.DataFrame`, we + // should generate `"pd.DataFrame"`. + return quote_annotation(parent_id, semantic, locator, stylist); + } + } + Some(Expr::Call(parent)) => { + if expr == parent.func.as_ref() { + // If we're quoting the function of a call, we need to quote the entire + // expression. For example, when quoting `DataFrame` in `DataFrame()`, we + // should generate `"DataFrame()"`. + return quote_annotation(parent_id, semantic, locator, stylist); + } + } + Some(Expr::BinOp(parent)) => { + if parent.op.is_bit_or() { + // If we're quoting the left or right side of a binary operation, we need to + // quote the entire expression. For example, when quoting `DataFrame` in + // `DataFrame | Series`, we should generate `"DataFrame | Series"`. + return quote_annotation(parent_id, semantic, locator, stylist); + } + } + _ => {} + } + } + + let annotation = locator.slice(expr); + + // If the annotation already contains a quote, avoid attempting to re-quote it. For example: + // ```python + // from typing import Literal + // + // Set[Literal["Foo"]] + // ``` + if annotation.contains('\'') || annotation.contains('"') { + return Err(anyhow::anyhow!("Annotation already contains a quote")); + } + + // If we're quoting a name, we need to quote the entire expression. + let quote = stylist.quote(); + let annotation = format!("{quote}{annotation}{quote}"); + Ok(Edit::range_replacement(annotation, expr.range())) +} diff --git a/crates/ruff_linter/src/rules/flake8_type_checking/imports.rs b/crates/ruff_linter/src/rules/flake8_type_checking/imports.rs new file mode 100644 index 0000000000000..92c65cf275e1c --- /dev/null +++ b/crates/ruff_linter/src/rules/flake8_type_checking/imports.rs @@ -0,0 +1,22 @@ +use ruff_python_semantic::{AnyImport, Binding, ResolvedReferenceId}; +use ruff_text_size::{Ranged, TextRange}; + +/// An import with its surrounding context. +pub(crate) struct ImportBinding<'a> { + /// The qualified name of the import (e.g., `typing.List` for `from typing import List`). + pub(crate) import: AnyImport<'a>, + /// The binding for the imported symbol. + pub(crate) binding: &'a Binding<'a>, + /// The first reference to the imported symbol. + pub(crate) reference_id: ResolvedReferenceId, + /// The trimmed range of the import (e.g., `List` in `from typing import List`). + pub(crate) range: TextRange, + /// The range of the import's parent statement. + pub(crate) parent_range: Option, +} + +impl Ranged for ImportBinding<'_> { + fn range(&self) -> TextRange { + self.range + } +} diff --git a/crates/ruff_linter/src/rules/flake8_type_checking/mod.rs b/crates/ruff_linter/src/rules/flake8_type_checking/mod.rs index 82b24755f4277..5e5a05cb6ff0e 100644 --- a/crates/ruff_linter/src/rules/flake8_type_checking/mod.rs +++ b/crates/ruff_linter/src/rules/flake8_type_checking/mod.rs @@ -1,5 +1,6 @@ //! Rules from [flake8-type-checking](https://pypi.org/project/flake8-type-checking/). pub(crate) mod helpers; +mod imports; pub(crate) mod rules; pub mod settings; @@ -33,10 +34,12 @@ mod tests { #[test_case(Rule::RuntimeImportInTypeCheckingBlock, Path::new("TCH004_7.py"))] #[test_case(Rule::RuntimeImportInTypeCheckingBlock, Path::new("TCH004_8.py"))] #[test_case(Rule::RuntimeImportInTypeCheckingBlock, Path::new("TCH004_9.py"))] + #[test_case(Rule::RuntimeImportInTypeCheckingBlock, Path::new("quote.py"))] #[test_case(Rule::TypingOnlyFirstPartyImport, Path::new("TCH001.py"))] #[test_case(Rule::TypingOnlyStandardLibraryImport, Path::new("TCH003.py"))] #[test_case(Rule::TypingOnlyStandardLibraryImport, Path::new("snapshot.py"))] #[test_case(Rule::TypingOnlyThirdPartyImport, Path::new("TCH002.py"))] + #[test_case(Rule::TypingOnlyThirdPartyImport, Path::new("quote.py"))] #[test_case(Rule::TypingOnlyThirdPartyImport, Path::new("singledispatch.py"))] #[test_case(Rule::TypingOnlyThirdPartyImport, Path::new("strict.py"))] #[test_case(Rule::TypingOnlyThirdPartyImport, Path::new("typing_modules_1.py"))] @@ -51,6 +54,24 @@ mod tests { Ok(()) } + #[test_case(Rule::RuntimeImportInTypeCheckingBlock, Path::new("quote.py"))] + #[test_case(Rule::TypingOnlyThirdPartyImport, Path::new("quote.py"))] + fn quote(rule_code: Rule, path: &Path) -> Result<()> { + let snapshot = format!("quote_{}_{}", rule_code.as_ref(), path.to_string_lossy()); + let diagnostics = test_path( + Path::new("flake8_type_checking").join(path).as_path(), + &settings::LinterSettings { + flake8_type_checking: super::settings::Settings { + quote_annotations: true, + ..Default::default() + }, + ..settings::LinterSettings::for_rule(rule_code) + }, + )?; + assert_messages!(snapshot, diagnostics); + Ok(()) + } + #[test_case(Rule::TypingOnlyThirdPartyImport, Path::new("strict.py"))] fn strict(rule_code: Rule, path: &Path) -> Result<()> { let diagnostics = test_path( @@ -109,7 +130,7 @@ mod tests { Path::new("flake8_type_checking").join(path).as_path(), &settings::LinterSettings { flake8_type_checking: super::settings::Settings { - runtime_evaluated_base_classes: vec![ + runtime_required_base_classes: vec![ "pydantic.BaseModel".to_string(), "sqlalchemy.orm.DeclarativeBase".to_string(), ], @@ -140,7 +161,7 @@ mod tests { Path::new("flake8_type_checking").join(path).as_path(), &settings::LinterSettings { flake8_type_checking: super::settings::Settings { - runtime_evaluated_decorators: vec![ + runtime_required_decorators: vec![ "attrs.define".to_string(), "attrs.frozen".to_string(), ], @@ -165,7 +186,7 @@ mod tests { Path::new("flake8_type_checking").join(path).as_path(), &settings::LinterSettings { flake8_type_checking: super::settings::Settings { - runtime_evaluated_base_classes: vec!["module.direct.MyBaseClass".to_string()], + runtime_required_base_classes: vec!["module.direct.MyBaseClass".to_string()], ..Default::default() }, ..settings::LinterSettings::for_rule(rule_code) diff --git a/crates/ruff_linter/src/rules/flake8_type_checking/rules/runtime_import_in_type_checking_block.rs b/crates/ruff_linter/src/rules/flake8_type_checking/rules/runtime_import_in_type_checking_block.rs index dea0f4007e0cd..5eee0365e12fa 100644 --- a/crates/ruff_linter/src/rules/flake8_type_checking/rules/runtime_import_in_type_checking_block.rs +++ b/crates/ruff_linter/src/rules/flake8_type_checking/rules/runtime_import_in_type_checking_block.rs @@ -5,13 +5,15 @@ use rustc_hash::FxHashMap; use ruff_diagnostics::{Diagnostic, Fix, FixAvailability, Violation}; use ruff_macros::{derive_message_formats, violation}; -use ruff_python_semantic::{AnyImport, Imported, NodeId, ResolvedReferenceId, Scope}; -use ruff_text_size::{Ranged, TextRange}; +use ruff_python_semantic::{Imported, NodeId, Scope}; +use ruff_text_size::Ranged; use crate::checkers::ast::Checker; use crate::codes::Rule; use crate::fix; use crate::importer::ImportedMembers; +use crate::rules::flake8_type_checking::helpers::quote_annotation; +use crate::rules::flake8_type_checking::imports::ImportBinding; /// ## What it does /// Checks for runtime imports defined in a type-checking block. @@ -20,6 +22,10 @@ use crate::importer::ImportedMembers; /// The type-checking block is not executed at runtime, so the import will not /// be available at runtime. /// +/// If [`flake8-type-checking.quote-annotations`] is set to `true`, +/// annotations will be wrapped in quotes if doing so would enable the +/// corresponding import to remain in the type-checking block. +/// /// ## Example /// ```python /// from typing import TYPE_CHECKING @@ -41,11 +47,15 @@ use crate::importer::ImportedMembers; /// foo.bar() /// ``` /// +/// ## Options +/// - `flake8-type-checking.quote-annotations` +/// /// ## References /// - [PEP 535](https://peps.python.org/pep-0563/#runtime-annotation-resolution-and-type-checking) #[violation] pub struct RuntimeImportInTypeCheckingBlock { qualified_name: String, + strategy: Strategy, } impl Violation for RuntimeImportInTypeCheckingBlock { @@ -53,17 +63,39 @@ impl Violation for RuntimeImportInTypeCheckingBlock { #[derive_message_formats] fn message(&self) -> String { - let RuntimeImportInTypeCheckingBlock { qualified_name } = self; - format!( - "Move import `{qualified_name}` out of type-checking block. Import is used for more than type hinting." - ) + let Self { + qualified_name, + strategy, + } = self; + match strategy { + Strategy::MoveImport => format!( + "Move import `{qualified_name}` out of type-checking block. Import is used for more than type hinting." + ), + Strategy::QuoteUsages => format!( + "Quote references to `{qualified_name}`. Import is in a type-checking block." + ), + } } fn fix_title(&self) -> Option { - Some("Move out of type-checking block".to_string()) + let Self { strategy, .. } = self; + match strategy { + Strategy::MoveImport => Some("Move out of type-checking block".to_string()), + Strategy::QuoteUsages => Some("Quote references".to_string()), + } } } +#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq)] +enum Action { + /// The import should be moved out of the type-checking block. + Move, + /// All usages of the import should be wrapped in quotes. + Quote, + /// The import should be ignored. + Ignore, +} + /// TCH004 pub(crate) fn runtime_import_in_type_checking_block( checker: &Checker, @@ -71,8 +103,7 @@ pub(crate) fn runtime_import_in_type_checking_block( diagnostics: &mut Vec, ) { // Collect all runtime imports by statement. - let mut errors_by_statement: FxHashMap> = FxHashMap::default(); - let mut ignores_by_statement: FxHashMap> = FxHashMap::default(); + let mut actions: FxHashMap<(NodeId, Action), Vec> = FxHashMap::default(); for binding_id in scope.binding_ids() { let binding = checker.semantic().binding(binding_id); @@ -101,6 +132,7 @@ pub(crate) fn runtime_import_in_type_checking_block( let import = ImportBinding { import, reference_id, + binding, range: binding.range(), parent_range: binding.parent_range(checker.semantic()), }; @@ -113,86 +145,153 @@ pub(crate) fn runtime_import_in_type_checking_block( ) }) { - ignores_by_statement - .entry(node_id) + actions + .entry((node_id, Action::Ignore)) .or_default() .push(import); } else { - errors_by_statement.entry(node_id).or_default().push(import); + // Determine whether the member should be fixed by moving the import out of the + // type-checking block, or by quoting its references. + if checker.settings.flake8_type_checking.quote_annotations + && binding.references().all(|reference_id| { + let reference = checker.semantic().reference(reference_id); + reference.context().is_typing() + || reference.in_runtime_evaluated_annotation() + }) + { + actions + .entry((node_id, Action::Quote)) + .or_default() + .push(import); + } else { + actions + .entry((node_id, Action::Move)) + .or_default() + .push(import); + } } } } - // Generate a diagnostic for every import, but share a fix across all imports within the same - // statement (excluding those that are ignored). - for (node_id, imports) in errors_by_statement { - let fix = fix_imports(checker, node_id, &imports).ok(); - - for ImportBinding { - import, - range, - parent_range, - .. - } in imports - { - let mut diagnostic = Diagnostic::new( - RuntimeImportInTypeCheckingBlock { - qualified_name: import.qualified_name(), - }, - range, - ); - if let Some(range) = parent_range { - diagnostic.set_parent(range.start()); + for ((node_id, action), imports) in actions { + match action { + // Generate a diagnostic for every import, but share a fix across all imports within the same + // statement (excluding those that are ignored). + Action::Move => { + let fix = move_imports(checker, node_id, &imports).ok(); + + for ImportBinding { + import, + range, + parent_range, + .. + } in imports + { + let mut diagnostic = Diagnostic::new( + RuntimeImportInTypeCheckingBlock { + qualified_name: import.qualified_name(), + strategy: Strategy::MoveImport, + }, + range, + ); + if let Some(range) = parent_range { + diagnostic.set_parent(range.start()); + } + if let Some(fix) = fix.as_ref() { + diagnostic.set_fix(fix.clone()); + } + diagnostics.push(diagnostic); + } } - if let Some(fix) = fix.as_ref() { - diagnostic.set_fix(fix.clone()); + + // Generate a diagnostic for every import, but share a fix across all imports within the same + // statement (excluding those that are ignored). + Action::Quote => { + let fix = quote_imports(checker, node_id, &imports).ok(); + + for ImportBinding { + import, + range, + parent_range, + .. + } in imports + { + let mut diagnostic = Diagnostic::new( + RuntimeImportInTypeCheckingBlock { + qualified_name: import.qualified_name(), + strategy: Strategy::QuoteUsages, + }, + range, + ); + if let Some(range) = parent_range { + diagnostic.set_parent(range.start()); + } + if let Some(fix) = fix.as_ref() { + diagnostic.set_fix(fix.clone()); + } + diagnostics.push(diagnostic); + } } - diagnostics.push(diagnostic); - } - } - // Separately, generate a diagnostic for every _ignored_ import, to ensure that the - // suppression comments aren't marked as unused. - for ImportBinding { - import, - range, - parent_range, - .. - } in ignores_by_statement.into_values().flatten() - { - let mut diagnostic = Diagnostic::new( - RuntimeImportInTypeCheckingBlock { - qualified_name: import.qualified_name(), - }, - range, - ); - if let Some(range) = parent_range { - diagnostic.set_parent(range.start()); + // Separately, generate a diagnostic for every _ignored_ import, to ensure that the + // suppression comments aren't marked as unused. + Action::Ignore => { + for ImportBinding { + import, + range, + parent_range, + .. + } in imports + { + let mut diagnostic = Diagnostic::new( + RuntimeImportInTypeCheckingBlock { + qualified_name: import.qualified_name(), + strategy: Strategy::MoveImport, + }, + range, + ); + if let Some(range) = parent_range { + diagnostic.set_parent(range.start()); + } + diagnostics.push(diagnostic); + } + } } - diagnostics.push(diagnostic); } } -/// A runtime-required import with its surrounding context. -struct ImportBinding<'a> { - /// The qualified name of the import (e.g., `typing.List` for `from typing import List`). - import: AnyImport<'a>, - /// The first reference to the imported symbol. - reference_id: ResolvedReferenceId, - /// The trimmed range of the import (e.g., `List` in `from typing import List`). - range: TextRange, - /// The range of the import's parent statement. - parent_range: Option, -} - -impl Ranged for ImportBinding<'_> { - fn range(&self) -> TextRange { - self.range - } +/// Generate a [`Fix`] to quote runtime usages for imports in a type-checking block. +fn quote_imports(checker: &Checker, node_id: NodeId, imports: &[ImportBinding]) -> Result { + let mut quote_reference_edits = imports + .iter() + .flat_map(|ImportBinding { binding, .. }| { + binding.references.iter().filter_map(|reference_id| { + let reference = checker.semantic().reference(*reference_id); + if reference.context().is_runtime() { + Some(quote_annotation( + reference.expression_id()?, + checker.semantic(), + checker.locator(), + checker.stylist(), + )) + } else { + None + } + }) + }) + .collect::>>()?; + let quote_reference_edit = quote_reference_edits + .pop() + .expect("Expected at least one reference"); + Ok( + Fix::unsafe_edits(quote_reference_edit, quote_reference_edits).isolate(Checker::isolation( + checker.semantic().parent_statement_id(node_id), + )), + ) } /// Generate a [`Fix`] to remove runtime imports from a type-checking block. -fn fix_imports(checker: &Checker, node_id: NodeId, imports: &[ImportBinding]) -> Result { +fn move_imports(checker: &Checker, node_id: NodeId, imports: &[ImportBinding]) -> Result { let statement = checker.semantic().statement(node_id); let parent = checker.semantic().parent_statement(node_id); @@ -236,3 +335,18 @@ fn fix_imports(checker: &Checker, node_id: NodeId, imports: &[ImportBinding]) -> ), ) } + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum Strategy { + /// The import should be moved out of the type-checking block. + /// + /// This is required when at least one reference to the symbol is in a runtime-required context. + /// For example, given `from foo import Bar`, `x = Bar()` would be runtime-required. + MoveImport, + /// All usages of the import should be wrapped in quotes. + /// + /// This is acceptable when all references to the symbol are in a runtime-evaluated, but not + /// runtime-required context. For example, given `from foo import Bar`, `x: Bar` would be + /// runtime-evaluated, but not runtime-required. + QuoteUsages, +} diff --git a/crates/ruff_linter/src/rules/flake8_type_checking/rules/typing_only_runtime_import.rs b/crates/ruff_linter/src/rules/flake8_type_checking/rules/typing_only_runtime_import.rs index 1ce6336d3b158..cb3adaed08368 100644 --- a/crates/ruff_linter/src/rules/flake8_type_checking/rules/typing_only_runtime_import.rs +++ b/crates/ruff_linter/src/rules/flake8_type_checking/rules/typing_only_runtime_import.rs @@ -5,13 +5,15 @@ use rustc_hash::FxHashMap; use ruff_diagnostics::{Diagnostic, DiagnosticKind, Fix, FixAvailability, Violation}; use ruff_macros::{derive_message_formats, violation}; -use ruff_python_semantic::{AnyImport, Binding, Imported, NodeId, ResolvedReferenceId, Scope}; -use ruff_text_size::{Ranged, TextRange}; +use ruff_python_semantic::{Binding, Imported, NodeId, Scope}; +use ruff_text_size::Ranged; use crate::checkers::ast::Checker; use crate::codes::Rule; use crate::fix; use crate::importer::ImportedMembers; +use crate::rules::flake8_type_checking::helpers::{is_typing_reference, quote_annotation}; +use crate::rules::flake8_type_checking::imports::ImportBinding; use crate::rules::isort::{categorize, ImportSection, ImportType}; /// ## What it does @@ -24,6 +26,10 @@ use crate::rules::isort::{categorize, ImportSection, ImportType}; /// instead be imported conditionally under an `if TYPE_CHECKING:` block to /// minimize runtime overhead. /// +/// If [`flake8-type-checking.quote-annotations`] is set to `true`, +/// annotations will be wrapped in quotes if doing so would enable the +/// corresponding import to be moved into an `if TYPE_CHECKING:` block. +/// /// If a class _requires_ that type annotations be available at runtime (as is /// the case for Pydantic, SQLAlchemy, and other libraries), consider using /// the [`flake8-type-checking.runtime-evaluated-base-classes`] and @@ -56,6 +62,7 @@ use crate::rules::isort::{categorize, ImportSection, ImportType}; /// ``` /// /// ## Options +/// - `flake8-type-checking.quote-annotations` /// - `flake8-type-checking.runtime-evaluated-base-classes` /// - `flake8-type-checking.runtime-evaluated-decorators` /// @@ -92,6 +99,10 @@ impl Violation for TypingOnlyFirstPartyImport { /// instead be imported conditionally under an `if TYPE_CHECKING:` block to /// minimize runtime overhead. /// +/// If [`flake8-type-checking.quote-annotations`] is set to `true`, +/// annotations will be wrapped in quotes if doing so would enable the +/// corresponding import to be moved into an `if TYPE_CHECKING:` block. +/// /// If a class _requires_ that type annotations be available at runtime (as is /// the case for Pydantic, SQLAlchemy, and other libraries), consider using /// the [`flake8-type-checking.runtime-evaluated-base-classes`] and @@ -124,6 +135,7 @@ impl Violation for TypingOnlyFirstPartyImport { /// ``` /// /// ## Options +/// - `flake8-type-checking.quote-annotations` /// - `flake8-type-checking.runtime-evaluated-base-classes` /// - `flake8-type-checking.runtime-evaluated-decorators` /// @@ -160,6 +172,10 @@ impl Violation for TypingOnlyThirdPartyImport { /// instead be imported conditionally under an `if TYPE_CHECKING:` block to /// minimize runtime overhead. /// +/// If [`flake8-type-checking.quote-annotations`] is set to `true`, +/// annotations will be wrapped in quotes if doing so would enable the +/// corresponding import to be moved into an `if TYPE_CHECKING:` block. +/// /// If a class _requires_ that type annotations be available at runtime (as is /// the case for Pydantic, SQLAlchemy, and other libraries), consider using /// the [`flake8-type-checking.runtime-evaluated-base-classes`] and @@ -192,6 +208,7 @@ impl Violation for TypingOnlyThirdPartyImport { /// ``` /// /// ## Options +/// - `flake8-type-checking.quote-annotations` /// - `flake8-type-checking.runtime-evaluated-base-classes` /// - `flake8-type-checking.runtime-evaluated-decorators` /// @@ -253,13 +270,12 @@ pub(crate) fn typing_only_runtime_import( }; if binding.context.is_runtime() - && binding.references().all(|reference_id| { - checker - .semantic() - .reference(reference_id) - .context() - .is_typing() - }) + && binding + .references() + .map(|reference_id| checker.semantic().reference(reference_id)) + .all(|reference| { + is_typing_reference(reference, &checker.settings.flake8_type_checking) + }) { let qualified_name = import.qualified_name(); @@ -310,6 +326,7 @@ pub(crate) fn typing_only_runtime_import( let import = ImportBinding { import, reference_id, + binding, range: binding.range(), parent_range: binding.parent_range(checker.semantic()), }; @@ -376,24 +393,6 @@ pub(crate) fn typing_only_runtime_import( } } -/// A runtime-required import with its surrounding context. -struct ImportBinding<'a> { - /// The qualified name of the import (e.g., `typing.List` for `from typing import List`). - import: AnyImport<'a>, - /// The first reference to the imported symbol. - reference_id: ResolvedReferenceId, - /// The trimmed range of the import (e.g., `List` in `from typing import List`). - range: TextRange, - /// The range of the import's parent statement. - parent_range: Option, -} - -impl Ranged for ImportBinding<'_> { - fn range(&self) -> TextRange { - self.range - } -} - /// Return the [`Rule`] for the given import type. fn rule_for(import_type: ImportType) -> Rule { match import_type { @@ -482,9 +481,34 @@ fn fix_imports(checker: &Checker, node_id: NodeId, imports: &[ImportBinding]) -> checker.source_type, )?; - Ok( - Fix::unsafe_edits(remove_import_edit, add_import_edit.into_edits()).isolate( - Checker::isolation(checker.semantic().parent_statement_id(node_id)), - ), + // Step 3) Quote any runtime usages of the referenced symbol. + let quote_reference_edits = imports + .iter() + .flat_map(|ImportBinding { binding, .. }| { + binding.references.iter().filter_map(|reference_id| { + let reference = checker.semantic().reference(*reference_id); + if reference.context().is_runtime() { + Some(quote_annotation( + reference.expression_id()?, + checker.semantic(), + checker.locator(), + checker.stylist(), + )) + } else { + None + } + }) + }) + .collect::>>()?; + + Ok(Fix::unsafe_edits( + remove_import_edit, + add_import_edit + .into_edits() + .into_iter() + .chain(quote_reference_edits), ) + .isolate(Checker::isolation( + checker.semantic().parent_statement_id(node_id), + ))) } diff --git a/crates/ruff_linter/src/rules/flake8_type_checking/settings.rs b/crates/ruff_linter/src/rules/flake8_type_checking/settings.rs index 425f02fe550a4..16baf1b91edbe 100644 --- a/crates/ruff_linter/src/rules/flake8_type_checking/settings.rs +++ b/crates/ruff_linter/src/rules/flake8_type_checking/settings.rs @@ -6,17 +6,19 @@ use ruff_macros::CacheKey; pub struct Settings { pub strict: bool, pub exempt_modules: Vec, - pub runtime_evaluated_base_classes: Vec, - pub runtime_evaluated_decorators: Vec, + pub runtime_required_base_classes: Vec, + pub runtime_required_decorators: Vec, + pub quote_annotations: bool, } impl Default for Settings { fn default() -> Self { Self { strict: false, - exempt_modules: vec!["typing".to_string()], - runtime_evaluated_base_classes: vec![], - runtime_evaluated_decorators: vec![], + exempt_modules: vec!["typing".to_string(), "typing_extensions".to_string()], + runtime_required_base_classes: vec![], + runtime_required_decorators: vec![], + quote_annotations: false, } } } diff --git a/crates/ruff_linter/src/rules/flake8_type_checking/snapshots/ruff_linter__rules__flake8_type_checking__tests__quote_runtime-import-in-type-checking-block_quote.py.snap b/crates/ruff_linter/src/rules/flake8_type_checking/snapshots/ruff_linter__rules__flake8_type_checking__tests__quote_runtime-import-in-type-checking-block_quote.py.snap new file mode 100644 index 0000000000000..0baeba9f62ec1 --- /dev/null +++ b/crates/ruff_linter/src/rules/flake8_type_checking/snapshots/ruff_linter__rules__flake8_type_checking__tests__quote_runtime-import-in-type-checking-block_quote.py.snap @@ -0,0 +1,22 @@ +--- +source: crates/ruff_linter/src/rules/flake8_type_checking/mod.rs +--- +quote.py:64:28: TCH004 [*] Quote references to `pandas.DataFrame`. Import is in a type-checking block. + | +63 | if TYPE_CHECKING: +64 | from pandas import DataFrame + | ^^^^^^^^^ TCH004 +65 | +66 | def func(value: DataFrame): + | + = help: Quote references + +ℹ Unsafe fix +63 63 | if TYPE_CHECKING: +64 64 | from pandas import DataFrame +65 65 | +66 |- def func(value: DataFrame): + 66 |+ def func(value: "DataFrame"): +67 67 | ... + + diff --git a/crates/ruff_linter/src/rules/flake8_type_checking/snapshots/ruff_linter__rules__flake8_type_checking__tests__quote_typing-only-third-party-import_quote.py.snap b/crates/ruff_linter/src/rules/flake8_type_checking/snapshots/ruff_linter__rules__flake8_type_checking__tests__quote_typing-only-third-party-import_quote.py.snap new file mode 100644 index 0000000000000..9c43070c0f5b8 --- /dev/null +++ b/crates/ruff_linter/src/rules/flake8_type_checking/snapshots/ruff_linter__rules__flake8_type_checking__tests__quote_typing-only-third-party-import_quote.py.snap @@ -0,0 +1,199 @@ +--- +source: crates/ruff_linter/src/rules/flake8_type_checking/mod.rs +--- +quote.py:2:24: TCH002 [*] Move third-party import `pandas.DataFrame` into a type-checking block + | +1 | def f(): +2 | from pandas import DataFrame + | ^^^^^^^^^ TCH002 +3 | +4 | def baz() -> DataFrame: + | + = help: Move into type-checking block + +ℹ Unsafe fix +1 |-def f(): + 1 |+from typing import TYPE_CHECKING + 2 |+ + 3 |+if TYPE_CHECKING: +2 4 | from pandas import DataFrame + 5 |+def f(): +3 6 | +4 |- def baz() -> DataFrame: + 7 |+ def baz() -> "DataFrame": +5 8 | ... +6 9 | +7 10 | + +quote.py:9:24: TCH002 [*] Move third-party import `pandas.DataFrame` into a type-checking block + | + 8 | def f(): + 9 | from pandas import DataFrame + | ^^^^^^^^^ TCH002 +10 | +11 | def baz() -> DataFrame[int]: + | + = help: Move into type-checking block + +ℹ Unsafe fix + 1 |+from typing import TYPE_CHECKING + 2 |+ + 3 |+if TYPE_CHECKING: + 4 |+ from pandas import DataFrame +1 5 | def f(): +2 6 | from pandas import DataFrame +3 7 | +-------------------------------------------------------------------------------- +6 10 | +7 11 | +8 12 | def f(): +9 |- from pandas import DataFrame +10 13 | +11 |- def baz() -> DataFrame[int]: + 14 |+ def baz() -> "DataFrame[int]": +12 15 | ... +13 16 | +14 17 | + +quote.py:16:24: TCH002 Move third-party import `pandas.DataFrame` into a type-checking block + | +15 | def f(): +16 | from pandas import DataFrame + | ^^^^^^^^^ TCH002 +17 | +18 | def baz() -> DataFrame["int"]: + | + = help: Move into type-checking block + +quote.py:23:22: TCH002 [*] Move third-party import `pandas` into a type-checking block + | +22 | def f(): +23 | import pandas as pd + | ^^ TCH002 +24 | +25 | def baz() -> pd.DataFrame: + | + = help: Move into type-checking block + +ℹ Unsafe fix + 1 |+from typing import TYPE_CHECKING + 2 |+ + 3 |+if TYPE_CHECKING: + 4 |+ import pandas as pd +1 5 | def f(): +2 6 | from pandas import DataFrame +3 7 | +-------------------------------------------------------------------------------- +20 24 | +21 25 | +22 26 | def f(): +23 |- import pandas as pd +24 27 | +25 |- def baz() -> pd.DataFrame: + 28 |+ def baz() -> "pd.DataFrame": +26 29 | ... +27 30 | +28 31 | + +quote.py:30:22: TCH002 [*] Move third-party import `pandas` into a type-checking block + | +29 | def f(): +30 | import pandas as pd + | ^^ TCH002 +31 | +32 | def baz() -> pd.DataFrame.Extra: + | + = help: Move into type-checking block + +ℹ Unsafe fix + 1 |+from typing import TYPE_CHECKING + 2 |+ + 3 |+if TYPE_CHECKING: + 4 |+ import pandas as pd +1 5 | def f(): +2 6 | from pandas import DataFrame +3 7 | +-------------------------------------------------------------------------------- +27 31 | +28 32 | +29 33 | def f(): +30 |- import pandas as pd +31 34 | +32 |- def baz() -> pd.DataFrame.Extra: + 35 |+ def baz() -> "pd.DataFrame.Extra": +33 36 | ... +34 37 | +35 38 | + +quote.py:37:22: TCH002 [*] Move third-party import `pandas` into a type-checking block + | +36 | def f(): +37 | import pandas as pd + | ^^ TCH002 +38 | +39 | def baz() -> pd.DataFrame | int: + | + = help: Move into type-checking block + +ℹ Unsafe fix + 1 |+from typing import TYPE_CHECKING + 2 |+ + 3 |+if TYPE_CHECKING: + 4 |+ import pandas as pd +1 5 | def f(): +2 6 | from pandas import DataFrame +3 7 | +-------------------------------------------------------------------------------- +34 38 | +35 39 | +36 40 | def f(): +37 |- import pandas as pd +38 41 | +39 |- def baz() -> pd.DataFrame | int: + 42 |+ def baz() -> "pd.DataFrame | int": +40 43 | ... +41 44 | +42 45 | + +quote.py:45:24: TCH002 [*] Move third-party import `pandas.DataFrame` into a type-checking block + | +44 | def f(): +45 | from pandas import DataFrame + | ^^^^^^^^^ TCH002 +46 | +47 | def baz() -> DataFrame(): + | + = help: Move into type-checking block + +ℹ Unsafe fix + 1 |+from typing import TYPE_CHECKING + 2 |+ + 3 |+if TYPE_CHECKING: + 4 |+ from pandas import DataFrame +1 5 | def f(): +2 6 | from pandas import DataFrame +3 7 | +-------------------------------------------------------------------------------- +42 46 | +43 47 | +44 48 | def f(): +45 |- from pandas import DataFrame +46 49 | +47 |- def baz() -> DataFrame(): + 50 |+ def baz() -> "DataFrame()": +48 51 | ... +49 52 | +50 53 | + +quote.py:54:24: TCH002 Move third-party import `pandas.DataFrame` into a type-checking block + | +52 | from typing import Literal +53 | +54 | from pandas import DataFrame + | ^^^^^^^^^ TCH002 +55 | +56 | def baz() -> DataFrame[Literal["int"]]: + | + = help: Move into type-checking block + + diff --git a/crates/ruff_linter/src/rules/flake8_type_checking/snapshots/ruff_linter__rules__flake8_type_checking__tests__runtime-import-in-type-checking-block_quote.py.snap b/crates/ruff_linter/src/rules/flake8_type_checking/snapshots/ruff_linter__rules__flake8_type_checking__tests__runtime-import-in-type-checking-block_quote.py.snap new file mode 100644 index 0000000000000..c09777853b97a --- /dev/null +++ b/crates/ruff_linter/src/rules/flake8_type_checking/snapshots/ruff_linter__rules__flake8_type_checking__tests__runtime-import-in-type-checking-block_quote.py.snap @@ -0,0 +1,29 @@ +--- +source: crates/ruff_linter/src/rules/flake8_type_checking/mod.rs +--- +quote.py:64:28: TCH004 [*] Move import `pandas.DataFrame` out of type-checking block. Import is used for more than type hinting. + | +63 | if TYPE_CHECKING: +64 | from pandas import DataFrame + | ^^^^^^^^^ TCH004 +65 | +66 | def func(value: DataFrame): + | + = help: Move out of type-checking block + +ℹ Unsafe fix + 1 |+from pandas import DataFrame +1 2 | def f(): +2 3 | from pandas import DataFrame +3 4 | +-------------------------------------------------------------------------------- +61 62 | from typing import TYPE_CHECKING +62 63 | +63 64 | if TYPE_CHECKING: +64 |- from pandas import DataFrame + 65 |+ pass +65 66 | +66 67 | def func(value: DataFrame): +67 68 | ... + + diff --git a/crates/ruff_linter/src/rules/flake8_type_checking/snapshots/ruff_linter__rules__flake8_type_checking__tests__typing-only-third-party-import_quote.py.snap b/crates/ruff_linter/src/rules/flake8_type_checking/snapshots/ruff_linter__rules__flake8_type_checking__tests__typing-only-third-party-import_quote.py.snap new file mode 100644 index 0000000000000..6c5ead27428ce --- /dev/null +++ b/crates/ruff_linter/src/rules/flake8_type_checking/snapshots/ruff_linter__rules__flake8_type_checking__tests__typing-only-third-party-import_quote.py.snap @@ -0,0 +1,4 @@ +--- +source: crates/ruff_linter/src/rules/flake8_type_checking/mod.rs +--- + diff --git a/crates/ruff_python_semantic/src/model.rs b/crates/ruff_python_semantic/src/model.rs index 82221f0f85dc8..d92dab84d0f5c 100644 --- a/crates/ruff_python_semantic/src/model.rs +++ b/crates/ruff_python_semantic/src/model.rs @@ -291,9 +291,12 @@ impl<'a> SemanticModel<'a> { if let Some(binding_id) = self.scopes.global().get(name.id.as_str()) { if !self.bindings[binding_id].is_unbound() { // Mark the binding as used. - let reference_id = - self.resolved_references - .push(ScopeId::global(), name.range, self.flags); + let reference_id = self.resolved_references.push( + ScopeId::global(), + self.node_id, + name.range, + self.flags, + ); self.bindings[binding_id].references.push(reference_id); // Mark any submodule aliases as used. @@ -302,6 +305,7 @@ impl<'a> SemanticModel<'a> { { let reference_id = self.resolved_references.push( ScopeId::global(), + self.node_id, name.range, self.flags, ); @@ -356,18 +360,24 @@ impl<'a> SemanticModel<'a> { if let Some(binding_id) = scope.get(name.id.as_str()) { // Mark the binding as used. - let reference_id = - self.resolved_references - .push(self.scope_id, name.range, self.flags); + let reference_id = self.resolved_references.push( + self.scope_id, + self.node_id, + name.range, + self.flags, + ); self.bindings[binding_id].references.push(reference_id); // Mark any submodule aliases as used. if let Some(binding_id) = self.resolve_submodule(name.id.as_str(), scope_id, binding_id) { - let reference_id = - self.resolved_references - .push(self.scope_id, name.range, self.flags); + let reference_id = self.resolved_references.push( + self.scope_id, + self.node_id, + name.range, + self.flags, + ); self.bindings[binding_id].references.push(reference_id); } @@ -431,9 +441,12 @@ impl<'a> SemanticModel<'a> { // The `x` in `print(x)` should resolve to the `x` in `x = 1`. BindingKind::UnboundException(Some(binding_id)) => { // Mark the binding as used. - let reference_id = - self.resolved_references - .push(self.scope_id, name.range, self.flags); + let reference_id = self.resolved_references.push( + self.scope_id, + self.node_id, + name.range, + self.flags, + ); self.bindings[binding_id].references.push(reference_id); // Mark any submodule aliases as used. @@ -442,6 +455,7 @@ impl<'a> SemanticModel<'a> { { let reference_id = self.resolved_references.push( self.scope_id, + self.node_id, name.range, self.flags, ); @@ -979,6 +993,23 @@ impl<'a> SemanticModel<'a> { &self.nodes[node_id] } + /// Given a [`Expr`], return its parent, if any. + #[inline] + pub fn parent_expression(&self, node_id: NodeId) -> Option<&'a Expr> { + self.nodes + .ancestor_ids(node_id) + .filter_map(|id| self.nodes[id].as_expression()) + .nth(1) + } + + /// Given a [`NodeId`], return the [`NodeId`] of the parent expression, if any. + pub fn parent_expression_id(&self, node_id: NodeId) -> Option { + self.nodes + .ancestor_ids(node_id) + .filter(|id| self.nodes[*id].is_expression()) + .nth(1) + } + /// Return the [`Stmt`] corresponding to the given [`NodeId`]. #[inline] pub fn statement(&self, node_id: NodeId) -> &'a Stmt { @@ -1007,11 +1038,10 @@ impl<'a> SemanticModel<'a> { /// Return the [`Expr`] corresponding to the given [`NodeId`]. #[inline] - pub fn expression(&self, node_id: NodeId) -> &'a Expr { + pub fn expression(&self, node_id: NodeId) -> Option<&'a Expr> { self.nodes .ancestor_ids(node_id) .find_map(|id| self.nodes[id].as_expression()) - .expect("No expression found") } /// Returns an [`Iterator`] over the expressions, starting from the given [`NodeId`]. @@ -1186,17 +1216,17 @@ impl<'a> SemanticModel<'a> { /// Add a reference to the given [`BindingId`] in the local scope. pub fn add_local_reference(&mut self, binding_id: BindingId, range: TextRange) { - let reference_id = self - .resolved_references - .push(self.scope_id, range, self.flags); + let reference_id = + self.resolved_references + .push(self.scope_id, self.node_id, range, self.flags); self.bindings[binding_id].references.push(reference_id); } /// Add a reference to the given [`BindingId`] in the global scope. pub fn add_global_reference(&mut self, binding_id: BindingId, range: TextRange) { - let reference_id = self - .resolved_references - .push(ScopeId::global(), range, self.flags); + let reference_id = + self.resolved_references + .push(ScopeId::global(), self.node_id, range, self.flags); self.bindings[binding_id].references.push(reference_id); } @@ -1299,10 +1329,16 @@ impl<'a> SemanticModel<'a> { .intersects(SemanticModelFlags::TYPING_ONLY_ANNOTATION) } - /// Return `true` if the model is in a runtime-required type annotation. - pub const fn in_runtime_annotation(&self) -> bool { + /// Return `true` if the context is in a runtime-evaluated type annotation. + pub const fn in_runtime_evaluated_annotation(&self) -> bool { + self.flags + .intersects(SemanticModelFlags::RUNTIME_EVALUATED_ANNOTATION) + } + + /// Return `true` if the context is in a runtime-required type annotation. + pub const fn in_runtime_required_annotation(&self) -> bool { self.flags - .intersects(SemanticModelFlags::RUNTIME_ANNOTATION) + .intersects(SemanticModelFlags::RUNTIME_REQUIRED_ANNOTATION) } /// Return `true` if the model is in a type definition. @@ -1474,8 +1510,9 @@ impl ShadowedBinding { bitflags! { /// Flags indicating the current model state. #[derive(Debug, Default, Copy, Clone, Eq, PartialEq)] - pub struct SemanticModelFlags: u16 { - /// The model is in a typing-time-only type annotation. + pub struct SemanticModelFlags: u32 { + /// The model is in a type annotation that will only be evaluated when running a type + /// checker. /// /// For example, the model could be visiting `int` in: /// ```python @@ -1490,7 +1527,7 @@ bitflags! { /// are any annotated assignments in module or class scopes. const TYPING_ONLY_ANNOTATION = 1 << 0; - /// The model is in a runtime type annotation. + /// The model is in a type annotation that will be evaluated at runtime. /// /// For example, the model could be visiting `int` in: /// ```python @@ -1504,7 +1541,27 @@ bitflags! { /// If `from __future__ import annotations` is used, all annotations are evaluated at /// typing time. Otherwise, all function argument annotations are evaluated at runtime, as /// are any annotated assignments in module or class scopes. - const RUNTIME_ANNOTATION = 1 << 1; + const RUNTIME_EVALUATED_ANNOTATION = 1 << 1; + + /// The model is in a type annotation that is _required_ to be available at runtime. + /// + /// For example, the context could be visiting `int` in: + /// ```python + /// from pydantic import BaseModel + /// + /// class Foo(BaseModel): + /// x: int + /// ``` + /// + /// In this case, Pydantic requires that the type annotation be available at runtime + /// in order to perform runtime type-checking. + /// + /// Unlike [`RUNTIME_EVALUATED_ANNOTATION`], annotations that are marked as + /// [`RUNTIME_REQUIRED_ANNOTATION`] cannot be deferred to typing time via conversion to a + /// forward reference (e.g., by wrapping the type in quotes), as the annotations are not + /// only required by the Python interpreter, but by runtime type checkers too. + const RUNTIME_REQUIRED_ANNOTATION = 1 << 2; + /// The model is in a type definition. /// @@ -1518,7 +1575,7 @@ bitflags! { /// All type annotations are also type definitions, but the converse is not true. /// In our example, `int` is a type definition but not a type annotation, as it /// doesn't appear in a type annotation context, but rather in a type definition. - const TYPE_DEFINITION = 1 << 2; + const TYPE_DEFINITION = 1 << 3; /// The model is in a (deferred) "simple" string type definition. /// @@ -1529,7 +1586,7 @@ bitflags! { /// /// "Simple" string type definitions are those that consist of a single string literal, /// as opposed to an implicitly concatenated string literal. - const SIMPLE_STRING_TYPE_DEFINITION = 1 << 3; + const SIMPLE_STRING_TYPE_DEFINITION = 1 << 4; /// The model is in a (deferred) "complex" string type definition. /// @@ -1540,7 +1597,7 @@ bitflags! { /// /// "Complex" string type definitions are those that consist of a implicitly concatenated /// string literals. These are uncommon but valid. - const COMPLEX_STRING_TYPE_DEFINITION = 1 << 4; + const COMPLEX_STRING_TYPE_DEFINITION = 1 << 5; /// The model is in a (deferred) `__future__` type definition. /// @@ -1553,7 +1610,7 @@ bitflags! { /// /// `__future__`-style type annotations are only enabled if the `annotations` feature /// is enabled via `from __future__ import annotations`. - const FUTURE_TYPE_DEFINITION = 1 << 5; + const FUTURE_TYPE_DEFINITION = 1 << 6; /// The model is in an exception handler. /// @@ -1564,7 +1621,7 @@ bitflags! { /// except Exception: /// x: int = 1 /// ``` - const EXCEPTION_HANDLER = 1 << 6; + const EXCEPTION_HANDLER = 1 << 7; /// The model is in an f-string. /// @@ -1572,7 +1629,7 @@ bitflags! { /// ```python /// f'{x}' /// ``` - const F_STRING = 1 << 7; + const F_STRING = 1 << 8; /// The model is in a boolean test. /// @@ -1584,7 +1641,7 @@ bitflags! { /// /// The implication is that the actual value returned by the current expression is /// not used, only its truthiness. - const BOOLEAN_TEST = 1 << 8; + const BOOLEAN_TEST = 1 << 9; /// The model is in a `typing::Literal` annotation. /// @@ -1593,7 +1650,7 @@ bitflags! { /// def f(x: Literal["A", "B", "C"]): /// ... /// ``` - const TYPING_LITERAL = 1 << 9; + const TYPING_LITERAL = 1 << 10; /// The model is in a subscript expression. /// @@ -1601,7 +1658,7 @@ bitflags! { /// ```python /// x["a"]["b"] /// ``` - const SUBSCRIPT = 1 << 10; + const SUBSCRIPT = 1 << 11; /// The model is in a type-checking block. /// @@ -1613,7 +1670,7 @@ bitflags! { /// if TYPE_CHECKING: /// x: int = 1 /// ``` - const TYPE_CHECKING_BLOCK = 1 << 11; + const TYPE_CHECKING_BLOCK = 1 << 12; /// The model has traversed past the "top-of-file" import boundary. /// @@ -1626,7 +1683,7 @@ bitflags! { /// /// x: int = 1 /// ``` - const IMPORT_BOUNDARY = 1 << 12; + const IMPORT_BOUNDARY = 1 << 13; /// The model has traversed past the `__future__` import boundary. /// @@ -1641,7 +1698,7 @@ bitflags! { /// /// Python considers it a syntax error to import from `__future__` after /// any other non-`__future__`-importing statements. - const FUTURES_BOUNDARY = 1 << 13; + const FUTURES_BOUNDARY = 1 << 14; /// `__future__`-style type annotations are enabled in this model. /// @@ -1653,7 +1710,7 @@ bitflags! { /// def f(x: int) -> int: /// ... /// ``` - const FUTURE_ANNOTATIONS = 1 << 14; + const FUTURE_ANNOTATIONS = 1 << 15; /// The model is in a type parameter definition. /// @@ -1663,10 +1720,11 @@ bitflags! { /// /// Record = TypeVar("Record") /// - const TYPE_PARAM_DEFINITION = 1 << 15; + const TYPE_PARAM_DEFINITION = 1 << 16; /// The context is in any type annotation. - const ANNOTATION = Self::TYPING_ONLY_ANNOTATION.bits() | Self::RUNTIME_ANNOTATION.bits(); + const ANNOTATION = Self::TYPING_ONLY_ANNOTATION.bits() | Self::RUNTIME_EVALUATED_ANNOTATION.bits() | Self::RUNTIME_REQUIRED_ANNOTATION.bits(); + /// The context is in any string type definition. const STRING_TYPE_DEFINITION = Self::SIMPLE_STRING_TYPE_DEFINITION.bits() diff --git a/crates/ruff_python_semantic/src/reference.rs b/crates/ruff_python_semantic/src/reference.rs index a963895b21e69..6bb807e1c8523 100644 --- a/crates/ruff_python_semantic/src/reference.rs +++ b/crates/ruff_python_semantic/src/reference.rs @@ -8,11 +8,14 @@ use ruff_text_size::{Ranged, TextRange}; use crate::context::ExecutionContext; use crate::scope::ScopeId; -use crate::{Exceptions, SemanticModelFlags}; +use crate::{Exceptions, NodeId, SemanticModelFlags}; /// A resolved read reference to a name in a program. #[derive(Debug, Clone)] pub struct ResolvedReference { + /// The expression that the reference occurs in. `None` if the reference is a global + /// reference or a reference via an augmented assignment. + node_id: Option, /// The scope in which the reference is defined. scope_id: ScopeId, /// The range of the reference in the source code. @@ -22,6 +25,11 @@ pub struct ResolvedReference { } impl ResolvedReference { + /// The expression that the reference occurs in. + pub const fn expression_id(&self) -> Option { + self.node_id + } + /// The scope in which the reference is defined. pub const fn scope_id(&self) -> ScopeId { self.scope_id @@ -35,6 +43,48 @@ impl ResolvedReference { ExecutionContext::Runtime } } + + /// Return `true` if the context is in a typing-only type annotation. + pub const fn in_typing_only_annotation(&self) -> bool { + self.flags + .intersects(SemanticModelFlags::TYPING_ONLY_ANNOTATION) + } + + /// Return `true` if the context is in a runtime-required type annotation. + pub const fn in_runtime_evaluated_annotation(&self) -> bool { + self.flags + .intersects(SemanticModelFlags::RUNTIME_EVALUATED_ANNOTATION) + } + + /// Return `true` if the context is in a "simple" string type definition. + pub const fn in_simple_string_type_definition(&self) -> bool { + self.flags + .intersects(SemanticModelFlags::SIMPLE_STRING_TYPE_DEFINITION) + } + + /// Return `true` if the context is in a "complex" string type definition. + pub const fn in_complex_string_type_definition(&self) -> bool { + self.flags + .intersects(SemanticModelFlags::COMPLEX_STRING_TYPE_DEFINITION) + } + + /// Return `true` if the context is in a `__future__` type definition. + pub const fn in_future_type_definition(&self) -> bool { + self.flags + .intersects(SemanticModelFlags::FUTURE_TYPE_DEFINITION) + } + + /// Return `true` if the context is in any kind of deferred type definition. + pub const fn in_deferred_type_definition(&self) -> bool { + self.flags + .intersects(SemanticModelFlags::DEFERRED_TYPE_DEFINITION) + } + + /// Return `true` if the context is in a type-checking block. + pub const fn in_type_checking_block(&self) -> bool { + self.flags + .intersects(SemanticModelFlags::TYPE_CHECKING_BLOCK) + } } impl Ranged for ResolvedReference { @@ -57,10 +107,12 @@ impl ResolvedReferences { pub(crate) fn push( &mut self, scope_id: ScopeId, + node_id: Option, range: TextRange, flags: SemanticModelFlags, ) -> ResolvedReferenceId { self.0.push(ResolvedReference { + node_id, scope_id, range, flags, diff --git a/crates/ruff_workspace/src/options.rs b/crates/ruff_workspace/src/options.rs index 10e5c6717696e..c0ce2c2e79f98 100644 --- a/crates/ruff_workspace/src/options.rs +++ b/crates/ruff_workspace/src/options.rs @@ -1642,6 +1642,57 @@ pub struct Flake8TypeCheckingOptions { "# )] pub runtime_evaluated_decorators: Option>, + + /// Whether to add quotes around type annotations, if doing so would allow + /// the corresponding import to be moved into a type-checking block. + /// + /// For example, in the following, Python requires that `Sequence` be + /// available at runtime, despite the fact that it's only used in a type + /// annotation: + /// + /// ```python + /// from collections.abc import Sequence + /// + /// + /// def func(value: Sequence[int]) -> None: + /// ... + /// ``` + /// + /// In other words, moving `from collections.abc import Sequence` into an + /// `if TYPE_CHECKING:` block above would cause a runtime error, as the + /// type would no longer be available at runtime. + /// + /// By default, Ruff will respect such runtime semantics and avoid moving + /// the import to prevent such runtime errors. + /// + /// Setting `quote-annotations` to `true` will instruct Ruff to add quotes + /// around the annotation (e.g., `"Sequence[int]"`), which in turn enables + /// Ruff to move the import into an `if TYPE_CHECKING:` block, like so: + /// + /// ```python + /// from typing import TYPE_CHECKING + /// + /// if TYPE_CHECKING: + /// from collections.abc import Sequence + /// + /// + /// def func(value: "Sequence[int]") -> None: + /// ... + /// ``` + /// + /// Note that this setting has no effect when `from __future__ import annotations` + /// is present, as `__future__` annotations are always treated equivalently + /// to quoted annotations. + #[option( + default = "false", + value_type = "bool", + example = r#" + # Add quotes around type annotations, if doing so would allow + # an import to be moved into a type-checking block. + quote-annotations = true + "# + )] + pub quote_annotations: Option, } impl Flake8TypeCheckingOptions { @@ -1651,8 +1702,9 @@ impl Flake8TypeCheckingOptions { exempt_modules: self .exempt_modules .unwrap_or_else(|| vec!["typing".to_string()]), - runtime_evaluated_base_classes: self.runtime_evaluated_base_classes.unwrap_or_default(), - runtime_evaluated_decorators: self.runtime_evaluated_decorators.unwrap_or_default(), + runtime_required_base_classes: self.runtime_evaluated_base_classes.unwrap_or_default(), + runtime_required_decorators: self.runtime_evaluated_decorators.unwrap_or_default(), + quote_annotations: self.quote_annotations.unwrap_or_default(), } } } diff --git a/ruff.schema.json b/ruff.schema.json index cbff63adcd807..8f67f9a12d787 100644 --- a/ruff.schema.json +++ b/ruff.schema.json @@ -1194,6 +1194,13 @@ "type": "string" } }, + "quote-annotations": { + "description": "Whether to add quotes around type annotations, if doing so would allow the corresponding import to be moved into a type-checking block.\n\nFor example, in the following, Python requires that `Sequence` be available at runtime, despite the fact that it's only used in a type annotation:\n\n```python from collections.abc import Sequence\n\ndef func(value: Sequence[int]) -> None: ... ```\n\nIn other words, moving `from collections.abc import Sequence` into an `if TYPE_CHECKING:` block above would cause a runtime error, as the type would no longer be available at runtime.\n\nBy default, Ruff will respect such runtime semantics and avoid moving the import to prevent such runtime errors.\n\nSetting `quote-annotations` to `true` will instruct Ruff to add quotes around the annotation (e.g., `\"Sequence[int]\"`), which in turn enables Ruff to move the import into an `if TYPE_CHECKING:` block, like so:\n\n```python from typing import TYPE_CHECKING\n\nif TYPE_CHECKING: from collections.abc import Sequence\n\ndef func(value: \"Sequence[int]\") -> None: ... ```\n\nNote that this setting has no effect when `from __future__ import annotations` is present, as `__future__` annotations are always treated equivalently to quoted annotations.", + "type": [ + "boolean", + "null" + ] + }, "runtime-evaluated-base-classes": { "description": "Exempt classes that list any of the enumerated classes as a base class from needing to be moved into type-checking blocks.\n\nCommon examples include Pydantic's `pydantic.BaseModel` and SQLAlchemy's `sqlalchemy.orm.DeclarativeBase`, but can also support user-defined classes that inherit from those base classes. For example, if you define a common `DeclarativeBase` subclass that's used throughout your project (e.g., `class Base(DeclarativeBase) ...` in `base.py`), you can add it to this list (`runtime-evaluated-base-classes = [\"base.Base\"]`) to exempt models from being moved into type-checking blocks.", "type": [