From 2f218f9ba26d5fb3f92fa311323b051378b92050 Mon Sep 17 00:00:00 2001 From: Alex Waygood Date: Wed, 24 Apr 2024 14:27:08 +0100 Subject: [PATCH 1/3] Refactor existing methods to determine if a class is an enumeration --- .../rules/flake8_pyi/rules/simple_defaults.rs | 63 +++---------------- .../rules/no_slots_in_str_subclass.rs | 21 ++----- .../ruff_python_semantic/src/analyze/class.rs | 13 ++++ 3 files changed, 25 insertions(+), 72 deletions(-) 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 70102ab892c9d..49ebe1be44d93 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,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,53 +470,6 @@ 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, - ) -> 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 @@ -676,21 +626,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; } } diff --git a/crates/ruff_linter/src/rules/flake8_slots/rules/no_slots_in_str_subclass.rs b/crates/ruff_linter/src/rules/flake8_slots/rules/no_slots_in_str_subclass.rs index a0ffe23e3b01f..68a562738d029 100644 --- a/crates/ruff_linter/src/rules/flake8_slots/rules/no_slots_in_str_subclass.rs +++ b/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" - ] - ) - }) -} diff --git a/crates/ruff_python_semantic/src/analyze/class.rs b/crates/ruff_python_semantic/src/analyze/class.rs index 5aaeebc9de8ad..44aa216d07da8 100644 --- a/crates/ruff_python_semantic/src/analyze/class.rs +++ b/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" + ] + ) + }) +} From 4f5a476213c508d5e3b0114e73d0040fa12dd5a8 Mon Sep 17 00:00:00 2001 From: Alex Waygood Date: Wed, 24 Apr 2024 14:40:25 +0100 Subject: [PATCH 2/3] Allow simple assignments to `None` in enum class scopes --- .../resources/test/fixtures/flake8_pyi/PYI026.py | 7 +++++++ .../resources/test/fixtures/flake8_pyi/PYI026.pyi | 7 +++++++ .../src/rules/flake8_pyi/rules/simple_defaults.rs | 14 +++++++++++--- 3 files changed, 25 insertions(+), 3 deletions(-) diff --git a/crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI026.py b/crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI026.py index afdbfe6480693..2fe996183cc00 100644 --- a/crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI026.py +++ b/crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI026.py @@ -15,5 +15,12 @@ AliasNone: typing.TypeAlias = None # these are ok +from enum import Enum + +class FooEnum(Enum): ... + +class BarEnum(FooEnum): + BAR = None + VarAlias = str AliasFoo = Foo diff --git a/crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI026.pyi b/crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI026.pyi index 87cb0ffe0d3d5..98bfa1bf49b4e 100644 --- a/crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI026.pyi +++ b/crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI026.pyi @@ -14,5 +14,12 @@ IntOrFloat: Foo = int | float AliasNone: typing.TypeAlias = None # these are ok +from enum import Enum + +class FooEnum(Enum): ... + +class BarEnum(FooEnum): + BAR = None + VarAlias = str AliasFoo = Foo 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 49ebe1be44d93..870e77e35a1ae 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 @@ -476,9 +476,17 @@ fn is_final_assignment(annotation: &Expr, value: &Expr, semantic: &SemanticModel /// 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 From 1142d5e4bb8a24500c33710c4f4b46de54681af4 Mon Sep 17 00:00:00 2001 From: Alex Waygood Date: Wed, 24 Apr 2024 14:45:06 +0100 Subject: [PATCH 3/3] Allow simple assignments to `None` in enum class scopes in stub files --- .../test/fixtures/flake8_pyi/PYI026.py | 4 +++- .../test/fixtures/flake8_pyi/PYI026.pyi | 3 +++ ..._flake8_pyi__tests__PYI026_PYI026.pyi.snap | 24 +++++++++++++++++++ ...e8_pyi__tests__py38_PYI026_PYI026.pyi.snap | 24 +++++++++++++++++++ 4 files changed, 54 insertions(+), 1 deletion(-) diff --git a/crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI026.py b/crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI026.py index 2fe996183cc00..16b91572a6c75 100644 --- a/crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI026.py +++ b/crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI026.py @@ -14,7 +14,9 @@ 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): ... diff --git a/crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI026.pyi b/crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI026.pyi index 98bfa1bf49b4e..21236c226b0a1 100644 --- a/crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI026.pyi +++ b/crates/ruff_linter/resources/test/fixtures/flake8_pyi/PYI026.pyi @@ -13,6 +13,9 @@ IntOrStr: TypeAlias = int | str IntOrFloat: Foo = int | float AliasNone: typing.TypeAlias = None +class NotAnEnum: + FLAG_THIS = None + # these are ok from enum import Enum diff --git a/crates/ruff_linter/src/rules/flake8_pyi/snapshots/ruff_linter__rules__flake8_pyi__tests__PYI026_PYI026.pyi.snap b/crates/ruff_linter/src/rules/flake8_pyi/snapshots/ruff_linter__rules__flake8_pyi__tests__PYI026_PYI026.pyi.snap index 80def65148c0c..91c4bbddd76ac 100644 --- a/crates/ruff_linter/src/rules/flake8_pyi/snapshots/ruff_linter__rules__flake8_pyi__tests__PYI026_PYI026.pyi.snap +++ b/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 diff --git a/crates/ruff_linter/src/rules/flake8_pyi/snapshots/ruff_linter__rules__flake8_pyi__tests__py38_PYI026_PYI026.pyi.snap b/crates/ruff_linter/src/rules/flake8_pyi/snapshots/ruff_linter__rules__flake8_pyi__tests__py38_PYI026_PYI026.pyi.snap index d10bf9016ef31..7c5cd3152327a 100644 --- a/crates/ruff_linter/src/rules/flake8_pyi/snapshots/ruff_linter__rules__flake8_pyi__tests__py38_PYI026_PYI026.pyi.snap +++ b/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