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

crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI026.py

@@ -14,6 +14,15 @@
 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

crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI026.pyi

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

crates/ruff_linter/src/rules/flake8_pyi/rules/simple_defaults.rs

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

crates/ruff_linter/src/rules/flake8_pyi/snapshots/ruff_linter__rules__flake8_pyi__tests__PYI026_PYI026.pyi.snap

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

crates/ruff_linter/src/rules/flake8_pyi/snapshots/ruff_linter__rules__flake8_pyi__tests__py38_PYI026_PYI026.pyi.snap

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

crates/ruff_linter/src/rules/flake8_slots/rules/no_slots_in_str_subclass.rs

@@ -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"
-            ]
-        )
-    })
-}

crates/ruff_python_semantic/src/analyze/class.rs

@@ -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"
+            ]
+        )
+    })
+}