diff --git a/crates/ruff_linter/resources/test/fixtures/flake8_type_checking/singledispatch.py b/crates/ruff_linter/resources/test/fixtures/flake8_type_checking/singledispatch.py new file mode 100644 index 0000000000000..a519a58b4f15b --- /dev/null +++ b/crates/ruff_linter/resources/test/fixtures/flake8_type_checking/singledispatch.py @@ -0,0 +1,34 @@ +"""Test module.""" +from __future__ import annotations + +from functools import singledispatch +from typing import TYPE_CHECKING + +from numpy import asarray +from numpy.typing import ArrayLike +from scipy.sparse import spmatrix +from pandas import DataFrame + +if TYPE_CHECKING: + from numpy import ndarray + + +@singledispatch +def to_array_or_mat(a: ArrayLike | spmatrix) -> ndarray | spmatrix: + """Convert arg to array or leaves it as sparse matrix.""" + msg = f"Unhandled type {type(a)}" + raise NotImplementedError(msg) + + +@to_array_or_mat.register +def _(a: ArrayLike) -> ndarray: + return asarray(a) + + +@to_array_or_mat.register +def _(a: spmatrix) -> spmatrix: + return a + + +def _(a: DataFrame) -> DataFrame: + return a diff --git a/crates/ruff_linter/src/checkers/ast/mod.rs b/crates/ruff_linter/src/checkers/ast/mod.rs index d24154badb6f7..11e8e704aa02f 100644 --- a/crates/ruff_linter/src/checkers/ast/mod.rs +++ b/crates/ruff_linter/src/checkers/ast/mod.rs @@ -492,6 +492,13 @@ where // are enabled. let runtime_annotation = !self.semantic.future_annotations(); + // The first parameter may be a single dispatch. + let mut singledispatch = + flake8_type_checking::helpers::is_singledispatch_implementation( + function_def, + self.semantic(), + ); + self.semantic.push_scope(ScopeKind::Type); if let Some(type_params) = type_params { @@ -505,7 +512,7 @@ where .chain(¶meters.kwonlyargs) { if let Some(expr) = ¶meter_with_default.parameter.annotation { - if runtime_annotation { + if runtime_annotation || singledispatch { self.visit_runtime_annotation(expr); } else { self.visit_annotation(expr); @@ -514,6 +521,7 @@ where if let Some(expr) = ¶meter_with_default.default { self.visit_expr(expr); } + singledispatch = false; } if let Some(arg) = ¶meters.vararg { if let Some(expr) = &arg.annotation { @@ -670,23 +678,24 @@ where // available at runtime. // See: https://docs.python.org/3/reference/simple_stmts.html#annotated-assignment-statements let runtime_annotation = if self.semantic.future_annotations() { - if self.semantic.current_scope().kind.is_class() { - let baseclasses = &self - .settings - .flake8_type_checking - .runtime_evaluated_base_classes; - let decorators = &self - .settings - .flake8_type_checking - .runtime_evaluated_decorators; - flake8_type_checking::helpers::runtime_evaluated( - baseclasses, - decorators, - &self.semantic, - ) - } else { - false - } + 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, 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 d07fe2d6cb1f2..0a51e151f4703 100644 --- a/crates/ruff_linter/src/rules/flake8_type_checking/helpers.rs +++ b/crates/ruff_linter/src/rules/flake8_type_checking/helpers.rs @@ -1,7 +1,7 @@ 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}; -use ruff_python_semantic::{Binding, BindingId, BindingKind, ScopeKind, SemanticModel}; +use ruff_python_ast::{self as ast, Expr}; +use ruff_python_semantic::{Binding, BindingId, BindingKind, SemanticModel}; use rustc_hash::FxHashSet; pub(crate) fn is_valid_runtime_import(binding: &Binding, semantic: &SemanticModel) -> bool { @@ -18,25 +18,26 @@ pub(crate) fn is_valid_runtime_import(binding: &Binding, semantic: &SemanticMode } } -pub(crate) fn runtime_evaluated( +pub(crate) fn runtime_evaluated_class( + class_def: &ast::StmtClassDef, base_classes: &[String], decorators: &[String], semantic: &SemanticModel, ) -> bool { - if !base_classes.is_empty() { - if runtime_evaluated_base_class(base_classes, semantic) { - return true; - } + if runtime_evaluated_base_class(class_def, base_classes, semantic) { + return true; } - if !decorators.is_empty() { - if runtime_evaluated_decorators(decorators, semantic) { - return true; - } + if runtime_evaluated_decorators(class_def, decorators, semantic) { + return true; } false } -fn runtime_evaluated_base_class(base_classes: &[String], semantic: &SemanticModel) -> bool { +fn runtime_evaluated_base_class( + class_def: &ast::StmtClassDef, + base_classes: &[String], + semantic: &SemanticModel, +) -> bool { fn inner( class_def: &ast::StmtClassDef, base_classes: &[String], @@ -78,19 +79,21 @@ fn runtime_evaluated_base_class(base_classes: &[String], semantic: &SemanticMode }) } - semantic - .current_scope() - .kind - .as_class() - .is_some_and(|class_def| { - inner(class_def, base_classes, semantic, &mut FxHashSet::default()) - }) + if base_classes.is_empty() { + return false; + } + + inner(class_def, base_classes, semantic, &mut FxHashSet::default()) } -fn runtime_evaluated_decorators(decorators: &[String], semantic: &SemanticModel) -> bool { - let ScopeKind::Class(class_def) = &semantic.current_scope().kind else { +fn runtime_evaluated_decorators( + class_def: &ast::StmtClassDef, + decorators: &[String], + semantic: &SemanticModel, +) -> bool { + if decorators.is_empty() { return false; - }; + } class_def.decorator_list.iter().any(|decorator| { semantic @@ -102,3 +105,72 @@ fn runtime_evaluated_decorators(decorators: &[String], semantic: &SemanticModel) }) }) } + +/// Returns `true` if a function is registered as a `singledispatch` interface. +/// +/// For example, `fun` below is a `singledispatch` interface: +/// ```python +/// from functools import singledispatch +/// +/// @singledispatch +/// def fun(arg, verbose=False): +/// ... +/// ``` +pub(crate) fn is_singledispatch_interface( + function_def: &ast::StmtFunctionDef, + semantic: &SemanticModel, +) -> bool { + function_def.decorator_list.iter().any(|decorator| { + semantic + .resolve_call_path(&decorator.expression) + .is_some_and(|call_path| { + matches!(call_path.as_slice(), ["functools", "singledispatch"]) + }) + }) +} + +/// Returns `true` if a function is registered as a `singledispatch` implementation. +/// +/// For example, `_` below is a `singledispatch` implementation: +/// For example: +/// ```python +/// from functools import singledispatch +/// +/// @singledispatch +/// def fun(arg, verbose=False): +/// ... +/// +/// @fun.register +/// def _(arg: int, verbose=False): +/// ... +/// ``` +pub(crate) fn is_singledispatch_implementation( + function_def: &ast::StmtFunctionDef, + semantic: &SemanticModel, +) -> bool { + function_def.decorator_list.iter().any(|decorator| { + let Expr::Attribute(attribute) = &decorator.expression else { + return false; + }; + + if attribute.attr.as_str() != "register" { + return false; + }; + + let Some(id) = semantic.lookup_attribute(attribute.value.as_ref()) else { + return false; + }; + + let binding = semantic.binding(id); + let Some(function_def) = binding + .kind + .as_function_definition() + .map(|id| &semantic.scopes[*id]) + .and_then(|scope| scope.kind.as_function()) + else { + return false; + }; + + is_singledispatch_interface(function_def, semantic) + }) +} 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 97e9ad7cd4af0..82b24755f4277 100644 --- a/crates/ruff_linter/src/rules/flake8_type_checking/mod.rs +++ b/crates/ruff_linter/src/rules/flake8_type_checking/mod.rs @@ -37,6 +37,7 @@ mod tests { #[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("singledispatch.py"))] #[test_case(Rule::TypingOnlyThirdPartyImport, Path::new("strict.py"))] #[test_case(Rule::TypingOnlyThirdPartyImport, Path::new("typing_modules_1.py"))] #[test_case(Rule::TypingOnlyThirdPartyImport, Path::new("typing_modules_2.py"))] diff --git a/crates/ruff_linter/src/rules/flake8_type_checking/snapshots/ruff_linter__rules__flake8_type_checking__tests__typing-only-third-party-import_singledispatch.py.snap b/crates/ruff_linter/src/rules/flake8_type_checking/snapshots/ruff_linter__rules__flake8_type_checking__tests__typing-only-third-party-import_singledispatch.py.snap new file mode 100644 index 0000000000000..5b646c6b34011 --- /dev/null +++ b/crates/ruff_linter/src/rules/flake8_type_checking/snapshots/ruff_linter__rules__flake8_type_checking__tests__typing-only-third-party-import_singledispatch.py.snap @@ -0,0 +1,27 @@ +--- +source: crates/ruff_linter/src/rules/flake8_type_checking/mod.rs +--- +singledispatch.py:10:20: TCH002 [*] Move third-party import `pandas.DataFrame` into a type-checking block + | + 8 | from numpy.typing import ArrayLike + 9 | from scipy.sparse import spmatrix +10 | from pandas import DataFrame + | ^^^^^^^^^ TCH002 +11 | +12 | if TYPE_CHECKING: + | + = help: Move into type-checking block + +ℹ Unsafe fix +7 7 | from numpy import asarray +8 8 | from numpy.typing import ArrayLike +9 9 | from scipy.sparse import spmatrix +10 |-from pandas import DataFrame +11 10 | +12 11 | if TYPE_CHECKING: + 12 |+ from pandas import DataFrame +13 13 | from numpy import ndarray +14 14 | +15 15 | + +