From 94178a03206fe2187c07820b6091f0e9dfd50404 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Sun, 19 Nov 2023 06:31:59 -0800 Subject: [PATCH] [`flake8-pyi`] Respect local enum subclasses in `simple-defaults` (`PYI052`) (#8767) We should reuse this approach in other rules, but this is a good start. Closes https://github.com/astral-sh/ruff/issues/8764. --- .../test/fixtures/flake8_pyi/PYI052.py | 14 ++++ .../test/fixtures/flake8_pyi/PYI052.pyi | 14 ++++ .../rules/flake8_pyi/rules/simple_defaults.rs | 69 +++++++++++++------ ..._flake8_pyi__tests__PYI052_PYI052.pyi.snap | 7 ++ 4 files changed, 84 insertions(+), 20 deletions(-) diff --git a/crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI052.py b/crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI052.py index 37a4f4d8671c5..7e740d3f86754 100644 --- a/crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI052.py +++ b/crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI052.py @@ -91,3 +91,17 @@ class Class1: field28 = builtins.str field29 = str field30 = str | bytes | None + +# We shouldn't emit Y052 for `enum` subclasses. +from enum import Enum + +class Foo(Enum): + FOO = 0 + BAR = 1 + +class Bar(Foo): + BAZ = 2 + BOP = 3 + +class Bop: + WIZ = 4 diff --git a/crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI052.pyi b/crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI052.pyi index 860ee255fb9b8..9e8237a4a4130 100644 --- a/crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI052.pyi +++ b/crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI052.pyi @@ -98,3 +98,17 @@ field27 = list[str] field28 = builtins.str field29 = str field30 = str | bytes | None + +# We shouldn't emit Y052 for `enum` subclasses. +from enum import Enum + +class Foo(Enum): + FOO = 0 + BAR = 1 + +class Bar(Foo): + BAZ = 2 + BOP = 3 + +class Bop: + WIZ = 4 diff --git a/crates/ruff_linter/src/rules/flake8_pyi/rules/simple_defaults.rs b/crates/ruff_linter/src/rules/flake8_pyi/rules/simple_defaults.rs index cd78602a7e717..a1fa2e58250a5 100644 --- a/crates/ruff_linter/src/rules/flake8_pyi/rules/simple_defaults.rs +++ b/crates/ruff_linter/src/rules/flake8_pyi/rules/simple_defaults.rs @@ -1,16 +1,17 @@ +use rustc_hash::FxHashSet; + use ruff_diagnostics::{AlwaysFixableViolation, Diagnostic, Edit, Fix, Violation}; use ruff_macros::{derive_message_formats, violation}; use ruff_python_ast::call_path::CallPath; use ruff_python_ast::{ self as ast, Arguments, Expr, Operator, ParameterWithDefault, Parameters, Stmt, UnaryOp, }; -use ruff_python_semantic::{ScopeKind, SemanticModel}; +use ruff_python_semantic::{BindingId, ScopeKind, SemanticModel}; use ruff_source_file::Locator; use ruff_text_size::Ranged; use crate::checkers::ast::Checker; use crate::importer::ImportRequest; - use crate::rules::flake8_pyi::rules::TypingModule; use crate::settings::types::PythonVersion; @@ -469,21 +470,51 @@ fn is_final_assignment(annotation: &Expr, value: &Expr, semantic: &SemanticModel } /// Returns `true` if the a class is an enum, based on its base classes. -fn is_enum(arguments: Option<&Arguments>, semantic: &SemanticModel) -> bool { - let Some(Arguments { args: bases, .. }) = arguments else { - return false; - }; - return bases.iter().any(|expr| { - semantic.resolve_call_path(expr).is_some_and(|call_path| { - matches!( - call_path.as_slice(), - [ - "enum", - "Enum" | "Flag" | "IntEnum" | "IntFlag" | "StrEnum" | "ReprEnum" - ] - ) +fn is_enum(class_def: &ast::StmtClassDef, semantic: &SemanticModel) -> bool { + fn inner( + class_def: &ast::StmtClassDef, + semantic: &SemanticModel, + seen: &mut FxHashSet, + ) -> bool { + let Some(Arguments { args: bases, .. }) = class_def.arguments.as_deref() else { + return false; + }; + + bases.iter().any(|expr| { + // If the base class is `enum.Enum`, `enum.Flag`, etc., then this is an enum. + if semantic.resolve_call_path(expr).is_some_and(|call_path| { + matches!( + call_path.as_slice(), + [ + "enum", + "Enum" | "Flag" | "IntEnum" | "IntFlag" | "StrEnum" | "ReprEnum" + ] + ) + }) { + return true; + } + + // If the base class extends `enum.Enum`, `enum.Flag`, etc., then this is an enum. + if let Some(id) = semantic.lookup_attribute(expr) { + if seen.insert(id) { + let binding = semantic.binding(id); + if let Some(base_class) = binding + .kind + .as_class_definition() + .map(|id| &semantic.scopes[*id]) + .and_then(|scope| scope.kind.as_class()) + { + if inner(base_class, semantic, seen) { + return true; + } + } + } + } + false }) - }); + } + + inner(class_def, semantic, &mut FxHashSet::default()) } /// Returns `true` if an [`Expr`] is a value that should be annotated with `typing.TypeAlias`. @@ -655,10 +686,8 @@ pub(crate) fn unannotated_assignment_in_stub( return; } - if let ScopeKind::Class(ast::StmtClassDef { arguments, .. }) = - checker.semantic().current_scope().kind - { - if is_enum(arguments.as_deref(), checker.semantic()) { + if let ScopeKind::Class(class_def) = checker.semantic().current_scope().kind { + if is_enum(class_def, checker.semantic()) { return; } } diff --git a/crates/ruff_linter/src/rules/flake8_pyi/snapshots/ruff_linter__rules__flake8_pyi__tests__PYI052_PYI052.pyi.snap b/crates/ruff_linter/src/rules/flake8_pyi/snapshots/ruff_linter__rules__flake8_pyi__tests__PYI052_PYI052.pyi.snap index 30f0ccedef4f2..0c39af3a9ff1e 100644 --- a/crates/ruff_linter/src/rules/flake8_pyi/snapshots/ruff_linter__rules__flake8_pyi__tests__PYI052_PYI052.pyi.snap +++ b/crates/ruff_linter/src/rules/flake8_pyi/snapshots/ruff_linter__rules__flake8_pyi__tests__PYI052_PYI052.pyi.snap @@ -131,4 +131,11 @@ PYI052.pyi:39:12: PYI052 Need type annotation for `field212` 41 | field22: Final = {"foo": 5} | +PYI052.pyi:114:11: PYI052 Need type annotation for `WIZ` + | +113 | class Bop: +114 | WIZ = 4 + | ^ PYI052 + | +