[`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.
This commit is contained in:
Charlie Marsh 2023-11-19 06:31:59 -08:00 committed by GitHub
parent 7165f8f05d
commit 94178a0320
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 84 additions and 20 deletions

View File

@ -91,3 +91,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

View File

@ -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

View File

@ -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<BindingId>,
) -> 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;
}
}

View File

@ -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
|