Skip to content

Commit

Permalink
[flake8-pyi] Respect local enum subclasses in simple-defaults (`P…
Browse files Browse the repository at this point in the history
…YI052`) (#8767)

We should reuse this approach in other rules, but this is a good start.

Closes #8764.
  • Loading branch information
charliermarsh authored Nov 19, 2023
1 parent 7165f8f commit 94178a0
Show file tree
Hide file tree
Showing 4 changed files with 84 additions and 20 deletions.
14 changes: 14 additions & 0 deletions crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI052.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
14 changes: 14 additions & 0 deletions crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI052.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -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
69 changes: 49 additions & 20 deletions crates/ruff_linter/src/rules/flake8_pyi/rules/simple_defaults.rs
Original file line number Diff line number Diff line change
@@ -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;

Expand Down Expand Up @@ -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`.
Expand Down Expand Up @@ -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;
}
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
|


0 comments on commit 94178a0

Please sign in to comment.