[flake8-pyi] Allow simple assignments to None in enum class scopes (PYI026) (#11128)

This commit is contained in:
Alex Waygood 2024-04-24 15:13:55 +01:00 committed by GitHub
parent 92814fd99b
commit 37af6e6147
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 104 additions and 76 deletions

View file

@ -14,6 +14,15 @@ IntOrStr: TypeAlias = int | str
IntOrFloat: Foo = int | float
AliasNone: typing.TypeAlias = None
# these are ok
class NotAnEnum:
NOT_A_STUB_SO_THIS_IS_FINE = None
from enum import Enum
class FooEnum(Enum): ...
class BarEnum(FooEnum):
BAR = None
VarAlias = str
AliasFoo = Foo

View file

@ -13,6 +13,16 @@ IntOrStr: TypeAlias = int | str
IntOrFloat: Foo = int | float
AliasNone: typing.TypeAlias = None
class NotAnEnum:
FLAG_THIS = None
# these are ok
from enum import Enum
class FooEnum(Enum): ...
class BarEnum(FooEnum):
BAR = None
VarAlias = str
AliasFoo = Foo

View file

@ -1,13 +1,10 @@
use rustc_hash::FxHashSet;
use ruff_diagnostics::{AlwaysFixableViolation, Diagnostic, Edit, Fix, Violation};
use ruff_macros::{derive_message_formats, violation};
use ruff_python_ast::helpers::map_subscript;
use ruff_python_ast::name::QualifiedName;
use ruff_python_ast::{
self as ast, Expr, Operator, ParameterWithDefault, Parameters, Stmt, UnaryOp,
};
use ruff_python_semantic::{BindingId, ScopeKind, SemanticModel};
use ruff_python_semantic::{analyze::class::is_enumeration, ScopeKind, SemanticModel};
use ruff_source_file::Locator;
use ruff_text_size::Ranged;
@ -473,62 +470,23 @@ fn is_final_assignment(annotation: &Expr, value: &Expr, semantic: &SemanticModel
false
}
/// Returns `true` if the a class is an enum, based on its base classes.
fn is_enum(class_def: &ast::StmtClassDef, semantic: &SemanticModel) -> bool {
fn inner(
class_def: &ast::StmtClassDef,
semantic: &SemanticModel,
seen: &mut FxHashSet<BindingId>,
) -> bool {
class_def.bases().iter().any(|expr| {
// If the base class is `enum.Enum`, `enum.Flag`, etc., then this is an enum.
if semantic
.resolve_qualified_name(map_subscript(expr))
.is_some_and(|qualified_name| {
matches!(
qualified_name.segments(),
[
"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(map_subscript(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`.
///
/// This is relatively conservative, as it's hard to reliably detect whether a right-hand side is a
/// valid type alias. In particular, this function checks for uses of `typing.Any`, `None`,
/// parameterized generics, and PEP 604-style unions.
fn is_annotatable_type_alias(value: &Expr, semantic: &SemanticModel) -> bool {
matches!(value, Expr::Subscript(_) | Expr::NoneLiteral(_))
|| is_valid_pep_604_union(value)
|| semantic.match_typing_expr(value, "Any")
if value.is_none_literal_expr() {
if let ScopeKind::Class(class_def) = semantic.current_scope().kind {
!is_enumeration(class_def, semantic)
} else {
true
}
} else {
value.is_subscript_expr()
|| is_valid_pep_604_union(value)
|| semantic.match_typing_expr(value, "Any")
}
}
/// PYI011
@ -676,21 +634,22 @@ pub(crate) fn unannotated_assignment_in_stub(
let Expr::Name(ast::ExprName { id, .. }) = target else {
return;
};
if is_special_assignment(target, checker.semantic()) {
let semantic = checker.semantic();
if is_special_assignment(target, semantic) {
return;
}
if is_type_var_like_call(value, checker.semantic()) {
if is_type_var_like_call(value, semantic) {
return;
}
if is_valid_default_value_without_annotation(value) {
return;
}
if !is_valid_default_value_with_annotation(value, true, checker.locator(), checker.semantic()) {
if !is_valid_default_value_with_annotation(value, true, checker.locator(), semantic) {
return;
}
if let ScopeKind::Class(class_def) = checker.semantic().current_scope().kind {
if is_enum(class_def, checker.semantic()) {
if let ScopeKind::Class(class_def) = semantic.current_scope().kind {
if is_enumeration(class_def, semantic) {
return;
}
}

View file

@ -114,4 +114,28 @@ PYI026.pyi:7:1: PYI026 [*] Use `typing.TypeAlias` for type alias, e.g., `AliasNo
9 9 | NewAny: typing.TypeAlias = Any
10 10 | OptionalStr: TypeAlias = typing.Optional[str]
PYI026.pyi:17:5: PYI026 [*] Use `typing.TypeAlias` for type alias, e.g., `FLAG_THIS: TypeAlias = None`
|
16 | class NotAnEnum:
17 | FLAG_THIS = None
| ^^^^^^^^^ PYI026
18 |
19 | # these are ok
|
= help: Add `TypeAlias` annotation
Safe fix
1 |-from typing import Literal, Any
1 |+from typing import Literal, Any, TypeAlias
2 2 |
3 3 | NewAny = Any
4 4 | OptionalStr = typing.Optional[str]
--------------------------------------------------------------------------------
14 14 | AliasNone: typing.TypeAlias = None
15 15 |
16 16 | class NotAnEnum:
17 |- FLAG_THIS = None
17 |+ FLAG_THIS: TypeAlias = None
18 18 |
19 19 | # these are ok
20 20 | from enum import Enum

View file

@ -114,4 +114,28 @@ PYI026.pyi:7:1: PYI026 [*] Use `typing_extensions.TypeAlias` for type alias, e.g
9 10 | NewAny: typing.TypeAlias = Any
10 11 | OptionalStr: TypeAlias = typing.Optional[str]
PYI026.pyi:17:5: PYI026 [*] Use `typing_extensions.TypeAlias` for type alias, e.g., `FLAG_THIS: TypeAlias = None`
|
16 | class NotAnEnum:
17 | FLAG_THIS = None
| ^^^^^^^^^ PYI026
18 |
19 | # these are ok
|
= help: Add `TypeAlias` annotation
Safe fix
1 1 | from typing import Literal, Any
2 |+import typing_extensions
2 3 |
3 4 | NewAny = Any
4 5 | OptionalStr = typing.Optional[str]
--------------------------------------------------------------------------------
14 15 | AliasNone: typing.TypeAlias = None
15 16 |
16 17 | class NotAnEnum:
17 |- FLAG_THIS = None
18 |+ FLAG_THIS: typing_extensions.TypeAlias = None
18 19 |
19 20 | # these are ok
20 21 | from enum import Enum

View file

@ -2,7 +2,7 @@ use ruff_diagnostics::{Diagnostic, Violation};
use ruff_macros::{derive_message_formats, violation};
use ruff_python_ast::identifier::Identifier;
use ruff_python_ast::{Arguments, Expr, Stmt, StmtClassDef};
use ruff_python_semantic::{analyze, SemanticModel};
use ruff_python_semantic::{analyze::class::is_enumeration, SemanticModel};
use crate::checkers::ast::Checker;
use crate::rules::flake8_slots::rules::helpers::has_slots;
@ -54,12 +54,14 @@ pub(crate) fn no_slots_in_str_subclass(checker: &mut Checker, stmt: &Stmt, class
return;
};
if !is_str_subclass(bases, checker.semantic()) {
let semantic = checker.semantic();
if !is_str_subclass(bases, semantic) {
return;
}
// Ignore subclasses of `enum.Enum` et al.
if is_enum_subclass(class, checker.semantic()) {
if is_enumeration(class, semantic) {
return;
}
@ -78,16 +80,3 @@ fn is_str_subclass(bases: &[Expr], semantic: &SemanticModel) -> bool {
.iter()
.any(|base| semantic.match_builtin_expr(base, "str"))
}
/// Returns `true` if the class is an enum subclass, at any depth.
fn is_enum_subclass(class_def: &StmtClassDef, semantic: &SemanticModel) -> bool {
analyze::class::any_qualified_name(class_def, semantic, &|qualified_name| {
matches!(
qualified_name.segments(),
[
"enum",
"Enum" | "IntEnum" | "StrEnum" | "Flag" | "IntFlag" | "ReprEnum" | "EnumCheck"
]
)
})
}

View file

@ -97,3 +97,16 @@ pub fn any_super_class(
inner(class_def, semantic, func, &mut FxHashSet::default())
}
/// Return `true` if `class_def` is a class that has one or more enum classes in its mro
pub fn is_enumeration(class_def: &ast::StmtClassDef, semantic: &SemanticModel) -> bool {
any_qualified_name(class_def, semantic, &|qualified_name| {
matches!(
qualified_name.segments(),
[
"enum",
"Enum" | "Flag" | "IntEnum" | "IntFlag" | "StrEnum" | "ReprEnum" | "CheckEnum"
]
)
})
}