[flake8-pyi] Implement PYI026

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

@@ -0,0 +1,19 @@
+import typing
+from typing import TypeAlias, Literal, Any
+
+NewAny = Any
+OptionalStr = typing.Optional[str]
+Foo = Literal["foo"]
+IntOrStr = int | str
+AliasNone = None
+
+NewAny: typing.TypeAlias = Any
+OptionalStr: TypeAlias = typing.Optional[str]
+Foo: typing.TypeAlias = Literal["foo"]
+IntOrStr: TypeAlias = int | str
+IntOrFloat: Foo = int | float
+AliasNone: typing.TypeAlias = None
+
+# these are ok
+VarAlias = str
+AliasFoo = Foo

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

@@ -0,0 +1,18 @@
+from typing import Literal, Any
+
+NewAny = Any
+OptionalStr = typing.Optional[str]
+Foo = Literal["foo"]
+IntOrStr = int | str
+AliasNone = None
+
+NewAny: typing.TypeAlias = Any
+OptionalStr: TypeAlias = typing.Optional[str]
+Foo: typing.TypeAlias = Literal["foo"]
+IntOrStr: TypeAlias = int | str
+IntOrFloat: Foo = int | float
+AliasNone: typing.TypeAlias = None
+
+# these are ok
+VarAlias = str
+AliasFoo = Foo

crates/ruff/src/checkers/ast/mod.rs

@@ -1557,6 +1557,7 @@ where
                         Rule::UnprefixedTypeParam,
                         Rule::AssignmentDefaultInStub,
                         Rule::UnannotatedAssignmentInStub,
+                        Rule::TypeAliasWithoutAnnotation,
                     ]) {
                         // Ignore assignments in function bodies; those are covered by other rules.
                         if !self
@@ -1575,6 +1576,11 @@ where
                                     self, targets, value,
                                 );
                             }
+                            if self.enabled(Rule::TypeAliasWithoutAnnotation) {
+                                flake8_pyi::rules::type_alias_without_annotation(
+                                    self, value, targets,
+                                );
+                            }
                         }
                     }
                 }

crates/ruff/src/codes.rs

@@ -633,6 +633,7 @@ pub fn code_to_rule(linter: Linter, code: &str) -> Option<(RuleGroup, Rule)> {
         (Flake8Pyi, "021") => (RuleGroup::Unspecified, rules::flake8_pyi::rules::DocstringInStub),
         (Flake8Pyi, "024") => (RuleGroup::Unspecified, rules::flake8_pyi::rules::CollectionsNamedTuple),
         (Flake8Pyi, "025") => (RuleGroup::Unspecified, rules::flake8_pyi::rules::UnaliasedCollectionsAbcSetImport),
+        (Flake8Pyi, "026") => (RuleGroup::Unspecified, rules::flake8_pyi::rules::TypeAliasWithoutAnnotation),
         (Flake8Pyi, "029") => (RuleGroup::Unspecified, rules::flake8_pyi::rules::StrOrReprDefinedInStub),
         (Flake8Pyi, "030") => (RuleGroup::Unspecified, rules::flake8_pyi::rules::UnnecessaryLiteralUnion),
         (Flake8Pyi, "032") => (RuleGroup::Unspecified, rules::flake8_pyi::rules::AnyEqNeAnnotation),

crates/ruff/src/rules/flake8_pyi/mod.rs

@@ -87,6 +87,8 @@ mod tests {
     #[test_case(Rule::UnrecognizedVersionInfoCheck, Path::new("PYI003.pyi"))]
     #[test_case(Rule::WrongTupleLengthVersionComparison, Path::new("PYI005.py"))]
     #[test_case(Rule::WrongTupleLengthVersionComparison, Path::new("PYI005.pyi"))]
+    #[test_case(Rule::TypeAliasWithoutAnnotation, Path::new("PYI026.py"))]
+    #[test_case(Rule::TypeAliasWithoutAnnotation, Path::new("PYI026.pyi"))]
     fn rules(rule_code: Rule, path: &Path) -> Result<()> {
         let snapshot = format!("{}_{}", rule_code.noqa_code(), path.to_string_lossy());
         let diagnostics = test_path(

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

@@ -9,6 +9,7 @@ use ruff_python_ast::source_code::Locator;
 use ruff_python_semantic::{ScopeKind, SemanticModel};
 
 use crate::checkers::ast::Checker;
+use crate::importer::ImportRequest;
 use crate::registry::AsRule;
 
 #[violation]
@@ -97,6 +98,47 @@ impl Violation for UnassignedSpecialVariableInStub {
     }
 }
 
+/// ## What it does
+/// Checks for type alias definitions that are not annotated with
+/// `typing.TypeAlias`.
+///
+/// ## Why is this bad?
+/// In Python, a type alias is defined by assigning a type to a variable (e.g.,
+/// `Vector = list[float]`).
+///
+/// It's best to annotate type aliases with the `typing.TypeAlias` type to
+/// make it clear that the statement is a type alias declaration, as opposed
+/// to a normal variable assignment.
+///
+/// ## Example
+/// ```python
+/// Vector = list[float]
+/// ```
+///
+/// Use instead:
+/// ```python
+/// from typing import TypeAlias
+///
+/// Vector: TypeAlias = list[float]
+/// ```
+#[violation]
+pub struct TypeAliasWithoutAnnotation {
+    name: String,
+    value: String,
+}
+
+impl AlwaysAutofixableViolation for TypeAliasWithoutAnnotation {
+    #[derive_message_formats]
+    fn message(&self) -> String {
+        let TypeAliasWithoutAnnotation { name, value } = self;
+        format!("Use `typing.TypeAlias` for type alias, e.g., `{name}: typing.TypeAlias = {value}`")
+    }
+
+    fn autofix_title(&self) -> String {
+        "Add `typing.TypeAlias` annotation".to_string()
+    }
+}
+
 fn is_allowed_negated_math_attribute(call_path: &CallPath) -> bool {
     matches!(call_path.as_slice(), ["math", "inf" | "e" | "pi" | "tau"])
 }
@@ -234,22 +276,39 @@ fn is_valid_default_value_with_annotation(
 
 /// Returns `true` if an [`Expr`] appears to be a valid PEP 604 union. (e.g. `int | None`)
 fn is_valid_pep_604_union(annotation: &Expr) -> bool {
-    match annotation {
-        Expr::BinOp(ast::ExprBinOp {
-            left,
-            op: Operator::BitOr,
-            right,
-            range: _,
-        }) => is_valid_pep_604_union(left) && is_valid_pep_604_union(right),
-        Expr::Name(_)
-        | Expr::Subscript(_)
-        | Expr::Attribute(_)
-        | Expr::Constant(ast::ExprConstant {
-            value: Constant::None,
-            ..
-        }) => true,
-        _ => false,
+    /// Returns `true` if an [`Expr`] appears to be a valid PEP 604 union member.
+    fn is_valid_pep_604_union_member(value: &Expr) -> bool {
+        match value {
+            Expr::BinOp(ast::ExprBinOp {
+                left,
+                op: Operator::BitOr,
+                right,
+                range: _,
+            }) => is_valid_pep_604_union_member(left) && is_valid_pep_604_union_member(right),
+            Expr::Name(_)
+            | Expr::Subscript(_)
+            | Expr::Attribute(_)
+            | Expr::Constant(ast::ExprConstant {
+                value: Constant::None,
+                ..
+            }) => true,
+            _ => false,
+        }
     }
+
+    // The top-level expression must be a bit-or operation.
+    let Expr::BinOp(ast::ExprBinOp {
+        left,
+        op: Operator::BitOr,
+        right,
+        range: _,
+    }) = annotation
+    else {
+        return false;
+    };
+
+    // The left and right operands must be valid union members.
+    is_valid_pep_604_union_member(left) && is_valid_pep_604_union_member(right)
 }
 
 /// Returns `true` if an [`Expr`] appears to be a valid default value without an annotation.
@@ -323,6 +382,23 @@ fn is_enum(bases: &[Expr], semantic: &SemanticModel) -> bool {
     });
 }
 
+/// 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::Constant(ast::ExprConstant {
+                value: Constant::None,
+                ..
+            }),
+    ) || is_valid_pep_604_union(value)
+        || semantic.match_typing_expr(value, "Any")
+}
+
 /// PYI011
 pub(crate) fn typed_argument_simple_defaults(checker: &mut Checker, arguments: &Arguments) {
     for ArgWithDefault {
@@ -523,3 +599,40 @@ pub(crate) fn unassigned_special_variable_in_stub(
         stmt.range(),
     ));
 }
+
+/// PIY026
+pub(crate) fn type_alias_without_annotation(checker: &mut Checker, value: &Expr, targets: &[Expr]) {
+    let [target] = targets else {
+        return;
+    };
+
+    let Expr::Name(ast::ExprName { id, .. }) = target else {
+        return;
+    };
+
+    if !is_annotatable_type_alias(value, checker.semantic()) {
+        return;
+    }
+
+    let mut diagnostic = Diagnostic::new(
+        TypeAliasWithoutAnnotation {
+            name: id.to_string(),
+            value: checker.generator().expr(value),
+        },
+        target.range(),
+    );
+    if checker.patch(diagnostic.kind.rule()) {
+        diagnostic.try_set_fix(|| {
+            let (import_edit, binding) = checker.importer.get_or_import_symbol(
+                &ImportRequest::import("typing", "TypeAlias"),
+                target.start(),
+                checker.semantic(),
+            )?;
+            Ok(Fix::suggested_edits(
+                Edit::range_replacement(format!("{id}: {binding}"), target.range()),
+                [import_edit],
+            ))
+        });
+    }
+    checker.diagnostics.push(diagnostic);
+}

crates/ruff/src/rules/flake8_pyi/snapshots/ruff__rules__flake8_pyi__tests__PYI026_PYI026.py.snap

@@ -0,0 +1,4 @@
+---
+source: crates/ruff/src/rules/flake8_pyi/mod.rs
+---
+

crates/ruff/src/rules/flake8_pyi/snapshots/ruff__rules__flake8_pyi__tests__PYI026_PYI026.pyi.snap

@@ -0,0 +1,117 @@
+---
+source: crates/ruff/src/rules/flake8_pyi/mod.rs
+---
+PYI026.pyi:3:1: PYI026 [*] Use `typing.TypeAlias` for type alias, e.g., `NewAny: typing.TypeAlias = Any`
+  |
+1 | from typing import Literal, Any
+2 | 
+3 | NewAny = Any
+  | ^^^^^^ PYI026
+4 | OptionalStr = typing.Optional[str]
+5 | Foo = Literal["foo"]
+  |
+  = help: Add `typing.TypeAlias` annotation
+
+ℹ Suggested fix
+1   |-from typing import Literal, Any
+  1 |+from typing import Literal, Any, TypeAlias
+2 2 | 
+3   |-NewAny = Any
+  3 |+NewAny: TypeAlias = Any
+4 4 | OptionalStr = typing.Optional[str]
+5 5 | Foo = Literal["foo"]
+6 6 | IntOrStr = int | str
+
+PYI026.pyi:4:1: PYI026 [*] Use `typing.TypeAlias` for type alias, e.g., `OptionalStr: typing.TypeAlias = typing.Optional[str]`
+  |
+3 | NewAny = Any
+4 | OptionalStr = typing.Optional[str]
+  | ^^^^^^^^^^^ PYI026
+5 | Foo = Literal["foo"]
+6 | IntOrStr = int | str
+  |
+  = help: Add `typing.TypeAlias` annotation
+
+ℹ Suggested fix
+1   |-from typing import Literal, Any
+  1 |+from typing import Literal, Any, TypeAlias
+2 2 | 
+3 3 | NewAny = Any
+4   |-OptionalStr = typing.Optional[str]
+  4 |+OptionalStr: TypeAlias = typing.Optional[str]
+5 5 | Foo = Literal["foo"]
+6 6 | IntOrStr = int | str
+7 7 | AliasNone = None
+
+PYI026.pyi:5:1: PYI026 [*] Use `typing.TypeAlias` for type alias, e.g., `Foo: typing.TypeAlias = Literal["foo"]`
+  |
+3 | NewAny = Any
+4 | OptionalStr = typing.Optional[str]
+5 | Foo = Literal["foo"]
+  | ^^^ PYI026
+6 | IntOrStr = int | str
+7 | AliasNone = None
+  |
+  = help: Add `typing.TypeAlias` annotation
+
+ℹ Suggested 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]
+5   |-Foo = Literal["foo"]
+  5 |+Foo: TypeAlias = Literal["foo"]
+6 6 | IntOrStr = int | str
+7 7 | AliasNone = None
+8 8 | 
+
+PYI026.pyi:6:1: PYI026 [*] Use `typing.TypeAlias` for type alias, e.g., `IntOrStr: typing.TypeAlias = int | str`
+  |
+4 | OptionalStr = typing.Optional[str]
+5 | Foo = Literal["foo"]
+6 | IntOrStr = int | str
+  | ^^^^^^^^ PYI026
+7 | AliasNone = None
+  |
+  = help: Add `typing.TypeAlias` annotation
+
+ℹ Suggested 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]
+5 5 | Foo = Literal["foo"]
+6   |-IntOrStr = int | str
+  6 |+IntOrStr: TypeAlias = int | str
+7 7 | AliasNone = None
+8 8 | 
+9 9 | NewAny: typing.TypeAlias = Any
+
+PYI026.pyi:7:1: PYI026 [*] Use `typing.TypeAlias` for type alias, e.g., `AliasNone: typing.TypeAlias = None`
+  |
+5 | Foo = Literal["foo"]
+6 | IntOrStr = int | str
+7 | AliasNone = None
+  | ^^^^^^^^^ PYI026
+8 | 
+9 | NewAny: typing.TypeAlias = Any
+  |
+  = help: Add `typing.TypeAlias` annotation
+
+ℹ Suggested 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]
+5 5 | Foo = Literal["foo"]
+6 6 | IntOrStr = int | str
+7   |-AliasNone = None
+  7 |+AliasNone: TypeAlias = None
+8 8 | 
+9 9 | NewAny: typing.TypeAlias = Any
+10 10 | OptionalStr: TypeAlias = typing.Optional[str]
+
+