[flake8-pyi] Implement PYI015 (#3728)

astral-sh · Mar 25, 2023 · 8ac7584 · 8ac7584
1 parent 4a1740a
commit 8ac7584
Show file tree

Hide file tree

Showing 11 changed files with 410 additions and 16 deletions.
diff --git a/crates/ruff/resources/test/fixtures/flake8_pyi/PYI015.py b/crates/ruff/resources/test/fixtures/flake8_pyi/PYI015.py
@@ -0,0 +1,48 @@
+import builtins
+import typing
+from typing import TypeAlias, Final
+
+field1: int
+field2: int = ...
+field3 = ...  # type: int  # Y033 Do not use type comments in stubs (e.g. use "x: int" instead of "x = ... # type: int")
+field4: int = 0
+field41: int = 0xFFFFFFFF
+field42: int = 1234567890
+field43: int = -0xFFFFFFFF
+field44: int = -1234567890
+field5 = 0  # type: int  # Y033 Do not use type comments in stubs (e.g. use "x: int" instead of "x = ... # type: int")  # Y052 Need type annotation for "field5"
+field6 = 0  # Y052 Need type annotation for "field6"
+field7 = b""  # Y052 Need type annotation for "field7"
+field71 = "foo"  # Y052 Need type annotation for "field71"
+field72: str = "foo"
+field8 = False  # Y052 Need type annotation for "field8"
+field81 = -1  # Y052 Need type annotation for "field81"
+field82: float = -98.43
+field83 = -42j  # Y052 Need type annotation for "field83"
+field84 = 5 + 42j  # Y052 Need type annotation for "field84"
+field85 = -5 - 42j  # Y052 Need type annotation for "field85"
+field9 = None  # Y026 Use typing_extensions.TypeAlias for type aliases, e.g. "field9: TypeAlias = None"
+Field95: TypeAlias = None
+Field96: TypeAlias = int | None
+Field97: TypeAlias = None | typing.SupportsInt | builtins.str | float | bool
+field19 = [1, 2, 3]  # Y052 Need type annotation for "field19"
+field191: list[int] = [1, 2, 3]
+field20 = (1, 2, 3)  # Y052 Need type annotation for "field20"
+field201: tuple[int, ...] = (1, 2, 3)
+field21 = {1, 2, 3}  # Y052 Need type annotation for "field21"
+field211: set[int] = {1, 2, 3}
+field212 = {"foo": "bar"}  # Y052 Need type annotation for "field212"
+field213: dict[str, str] = {"foo": "bar"}
+field22: Final = {"foo": 5}
+field221: list[int] = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]  # Y015 Only simple default values are allowed for assignments
+field223: list[int] = [*range(10)]  # Y015 Only simple default values are allowed for assignments
+field224: list[int] = list(range(10))  # Y015 Only simple default values are allowed for assignments
+field225: list[object] = [{}, 1, 2]  # Y015 Only simple default values are allowed for assignments
+field226: tuple[str | tuple[str, ...], ...] = ("foo", ("foo", "bar"))  # Y015 Only simple default values are allowed for assignments
+field227: dict[str, object] = {"foo": {"foo": "bar"}}  # Y015 Only simple default values are allowed for assignments
+field228: dict[str, list[object]] = {"foo": []}  # Y015 Only simple default values are allowed for assignments
+# When parsed, this case results in `None` being placed in the `.keys` list for the `ast.Dict` node
+field229: dict[int, int] = {1: 2, **{3: 4}}  # Y015 Only simple default values are allowed for assignments
+field23 = "foo" + "bar"  # Y015 Only simple default values are allowed for assignments
+field24 = b"foo" + b"bar"  # Y015 Only simple default values are allowed for assignments
+field25 = 5 * 5  # Y015 Only simple default values are allowed for assignments
diff --git a/crates/ruff/resources/test/fixtures/flake8_pyi/PYI015.pyi b/crates/ruff/resources/test/fixtures/flake8_pyi/PYI015.pyi
@@ -0,0 +1,51 @@
+import builtins
+import typing
+from typing import TypeAlias, Final
+
+# We shouldn't emit Y015 for simple default values
+field1: int
+field2: int = ...
+field3 = ...  # type: int  # Y033 Do not use type comments in stubs (e.g. use "x: int" instead of "x = ... # type: int")
+field4: int = 0
+field41: int = 0xFFFFFFFF
+field42: int = 1234567890
+field43: int = -0xFFFFFFFF
+field44: int = -1234567890
+field5 = 0  # type: int  # Y033 Do not use type comments in stubs (e.g. use "x: int" instead of "x = ... # type: int")  # Y052 Need type annotation for "field5"
+field6 = 0  # Y052 Need type annotation for "field6"
+field7 = b""  # Y052 Need type annotation for "field7"
+field71 = "foo"  # Y052 Need type annotation for "field71"
+field72: str = "foo"
+field8 = False  # Y052 Need type annotation for "field8"
+field81 = -1  # Y052 Need type annotation for "field81"
+field82: float = -98.43
+field83 = -42j  # Y052 Need type annotation for "field83"
+field84 = 5 + 42j  # Y052 Need type annotation for "field84"
+field85 = -5 - 42j  # Y052 Need type annotation for "field85"
+field9 = None  # Y026 Use typing_extensions.TypeAlias for type aliases, e.g. "field9: TypeAlias = None"
+Field95: TypeAlias = None
+Field96: TypeAlias = int | None
+Field97: TypeAlias = None | typing.SupportsInt | builtins.str | float | bool
+field19 = [1, 2, 3]  # Y052 Need type annotation for "field19"
+field191: list[int] = [1, 2, 3]
+field20 = (1, 2, 3)  # Y052 Need type annotation for "field20"
+field201: tuple[int, ...] = (1, 2, 3)
+field21 = {1, 2, 3}  # Y052 Need type annotation for "field21"
+field211: set[int] = {1, 2, 3}
+field212 = {"foo": "bar"}  # Y052 Need type annotation for "field212"
+field213: dict[str, str] = {"foo": "bar"}
+field22: Final = {"foo": 5}
+
+# We *should* emit Y015 for more complex default values
+field221: list[int] = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]  # Y015 Only simple default values are allowed for assignments
+field223: list[int] = [*range(10)]  # Y015 Only simple default values are allowed for assignments
+field224: list[int] = list(range(10))  # Y015 Only simple default values are allowed for assignments
+field225: list[object] = [{}, 1, 2]  # Y015 Only simple default values are allowed for assignments
+field226: tuple[str | tuple[str, ...], ...] = ("foo", ("foo", "bar"))  # Y015 Only simple default values are allowed for assignments
+field227: dict[str, object] = {"foo": {"foo": "bar"}}  # Y015 Only simple default values are allowed for assignments
+field228: dict[str, list[object]] = {"foo": []}  # Y015 Only simple default values are allowed for assignments
+# When parsed, this case results in `None` being placed in the `.keys` list for the `ast.Dict` node
+field229: dict[int, int] = {1: 2, **{3: 4}}  # Y015 Only simple default values are allowed for assignments
+field23 = "foo" + "bar"  # Y015 Only simple default values are allowed for assignments
+field24 = b"foo" + b"bar"  # Y015 Only simple default values are allowed for assignments
+field25 = 5 * 5  # Y015 Only simple default values are allowed for assignments
diff --git a/crates/ruff/src/checkers/ast/mod.rs b/crates/ruff/src/checkers/ast/mod.rs
@@ -1801,8 +1801,19 @@ where
                         self.diagnostics.push(diagnostic);
                     }
                 }
+
+                if self.is_stub {
+                    if self.settings.rules.enabled(Rule::AssignmentDefaultInStub) {
+                        flake8_pyi::rules::assignment_default_in_stub(self, value, None);
+                    }
+                }
             }
-            StmtKind::AnnAssign { target, value, .. } => {
+            StmtKind::AnnAssign {
+                target,
+                value,
+                annotation,
+                ..
+            } => {
                 if self.settings.rules.enabled(Rule::LambdaAssignment) {
                     if let Some(value) = value {
                         pycodestyle::rules::lambda_assignment(self, target, value, stmt);
@@ -1820,6 +1831,17 @@ where
                         stmt,
                     );
                 }
+                if self.is_stub {
+                    if let Some(value) = value {
+                        if self.settings.rules.enabled(Rule::AssignmentDefaultInStub) {
+                            flake8_pyi::rules::assignment_default_in_stub(
+                                self,
+                                value,
+                                Some(annotation),
+                            );
+                        }
+                    }
+                }
             }
             StmtKind::Delete { targets } => {
                 if self.settings.rules.enabled(Rule::GlobalStatement) {

diff --git a/crates/ruff/src/codes.rs b/crates/ruff/src/codes.rs
@@ -565,6 +565,7 @@ pub fn code_to_rule(linter: Linter, code: &str) -> Option<Rule> {
         (Flake8Pyi, "010") => Rule::NonEmptyStubBody,
         (Flake8Pyi, "011") => Rule::TypedArgumentDefaultInStub,
         (Flake8Pyi, "014") => Rule::ArgumentDefaultInStub,
+        (Flake8Pyi, "015") => Rule::AssignmentDefaultInStub,
         (Flake8Pyi, "021") => Rule::DocstringInStub,
         (Flake8Pyi, "033") => Rule::TypeCommentInStub,
 

diff --git a/crates/ruff/src/registry.rs b/crates/ruff/src/registry.rs
@@ -514,16 +514,17 @@ ruff_macros::register_rules!(
     rules::flake8_errmsg::rules::FStringInException,
     rules::flake8_errmsg::rules::DotFormatInException,
     // flake8-pyi
-    rules::flake8_pyi::rules::UnprefixedTypeParam,
+    rules::flake8_pyi::rules::ArgumentDefaultInStub,
+    rules::flake8_pyi::rules::AssignmentDefaultInStub,
     rules::flake8_pyi::rules::BadVersionInfoComparison,
-    rules::flake8_pyi::rules::UnrecognizedPlatformCheck,
-    rules::flake8_pyi::rules::UnrecognizedPlatformName,
-    rules::flake8_pyi::rules::PassStatementStubBody,
-    rules::flake8_pyi::rules::NonEmptyStubBody,
     rules::flake8_pyi::rules::DocstringInStub,
-    rules::flake8_pyi::rules::TypedArgumentDefaultInStub,
-    rules::flake8_pyi::rules::ArgumentDefaultInStub,
+    rules::flake8_pyi::rules::NonEmptyStubBody,
+    rules::flake8_pyi::rules::PassStatementStubBody,
     rules::flake8_pyi::rules::TypeCommentInStub,
+    rules::flake8_pyi::rules::TypedArgumentDefaultInStub,
+    rules::flake8_pyi::rules::UnprefixedTypeParam,
+    rules::flake8_pyi::rules::UnrecognizedPlatformCheck,
+    rules::flake8_pyi::rules::UnrecognizedPlatformName,
     // flake8-pytest-style
     rules::flake8_pytest_style::rules::PytestFixtureIncorrectParenthesesStyle,
     rules::flake8_pytest_style::rules::PytestFixturePositionalArgs,

diff --git a/crates/ruff/src/rules/flake8_pyi/mod.rs b/crates/ruff/src/rules/flake8_pyi/mod.rs
@@ -13,22 +13,24 @@ mod tests {
     use crate::settings;
     use crate::test::test_path;
 
-    #[test_case(Rule::UnprefixedTypeParam, Path::new("PYI001.pyi"))]
     #[test_case(Rule::UnprefixedTypeParam, Path::new("PYI001.py"))]
-    #[test_case(Rule::BadVersionInfoComparison, Path::new("PYI006.pyi"))]
+    #[test_case(Rule::UnprefixedTypeParam, Path::new("PYI001.pyi"))]
     #[test_case(Rule::BadVersionInfoComparison, Path::new("PYI006.py"))]
-    #[test_case(Rule::UnrecognizedPlatformCheck, Path::new("PYI007.pyi"))]
+    #[test_case(Rule::BadVersionInfoComparison, Path::new("PYI006.pyi"))]
     #[test_case(Rule::UnrecognizedPlatformCheck, Path::new("PYI007.py"))]
-    #[test_case(Rule::UnrecognizedPlatformName, Path::new("PYI008.pyi"))]
+    #[test_case(Rule::UnrecognizedPlatformCheck, Path::new("PYI007.pyi"))]
     #[test_case(Rule::UnrecognizedPlatformName, Path::new("PYI008.py"))]
-    #[test_case(Rule::NonEmptyStubBody, Path::new("PYI010.py"))]
-    #[test_case(Rule::NonEmptyStubBody, Path::new("PYI010.pyi"))]
+    #[test_case(Rule::UnrecognizedPlatformName, Path::new("PYI008.pyi"))]
     #[test_case(Rule::PassStatementStubBody, Path::new("PYI009.py"))]
     #[test_case(Rule::PassStatementStubBody, Path::new("PYI009.pyi"))]
+    #[test_case(Rule::NonEmptyStubBody, Path::new("PYI010.py"))]
+    #[test_case(Rule::NonEmptyStubBody, Path::new("PYI010.pyi"))]
     #[test_case(Rule::TypedArgumentDefaultInStub, Path::new("PYI011.py"))]
     #[test_case(Rule::TypedArgumentDefaultInStub, Path::new("PYI011.pyi"))]
     #[test_case(Rule::ArgumentDefaultInStub, Path::new("PYI014.py"))]
     #[test_case(Rule::ArgumentDefaultInStub, Path::new("PYI014.pyi"))]
+    #[test_case(Rule::AssignmentDefaultInStub, Path::new("PYI015.py"))]
+    #[test_case(Rule::AssignmentDefaultInStub, Path::new("PYI015.pyi"))]
     #[test_case(Rule::DocstringInStub, Path::new("PYI021.py"))]
     #[test_case(Rule::DocstringInStub, Path::new("PYI021.pyi"))]
     #[test_case(Rule::TypeCommentInStub, Path::new("PYI033.py"))]

diff --git a/crates/ruff/src/rules/flake8_pyi/rules/mod.rs b/crates/ruff/src/rules/flake8_pyi/rules/mod.rs
@@ -4,8 +4,8 @@ pub use non_empty_stub_body::{non_empty_stub_body, NonEmptyStubBody};
 pub use pass_statement_stub_body::{pass_statement_stub_body, PassStatementStubBody};
 pub use prefix_type_params::{prefix_type_params, UnprefixedTypeParam};
 pub use simple_defaults::{
-    argument_simple_defaults, typed_argument_simple_defaults, ArgumentDefaultInStub,
-    TypedArgumentDefaultInStub,
+    argument_simple_defaults, assignment_default_in_stub, typed_argument_simple_defaults,
+    ArgumentDefaultInStub, AssignmentDefaultInStub, TypedArgumentDefaultInStub,
 };
 pub use type_comment_in_stub::{type_comment_in_stub, TypeCommentInStub};
 pub use unrecognized_platform::{

diff --git a/crates/ruff/src/rules/flake8_pyi/rules/simple_defaults.rs b/crates/ruff/src/rules/flake8_pyi/rules/simple_defaults.rs
@@ -37,6 +37,21 @@ impl AlwaysAutofixableViolation for ArgumentDefaultInStub {
     }
 }
 
+#[violation]
+pub struct AssignmentDefaultInStub;
+
+/// PYI015
+impl AlwaysAutofixableViolation for AssignmentDefaultInStub {
+    #[derive_message_formats]
+    fn message(&self) -> String {
+        format!("Only simple default values allowed for assignments")
+    }
+
+    fn autofix_title(&self) -> String {
+        "Replace default value with `...`".to_string()
+    }
+}
+
 const ALLOWED_MATH_ATTRIBUTES_IN_DEFAULTS: &[&[&str]] = &[
     &["math", "inf"],
     &["math", "nan"],
@@ -319,3 +334,25 @@ pub fn argument_simple_defaults(checker: &mut Checker, args: &Arguments) {
         }
     }
 }
+
+/// PYI015
+pub fn assignment_default_in_stub(checker: &mut Checker, value: &Expr, annotation: Option<&Expr>) {
+    if annotation.map_or(false, |annotation| {
+        checker.ctx.match_typing_expr(annotation, "TypeAlias")
+    }) {
+        return;
+    }
+    if !is_valid_default_value_with_annotation(value, checker, true) {
+        let mut diagnostic = Diagnostic::new(AssignmentDefaultInStub, Range::from(value));
+
+        if checker.patch(diagnostic.kind.rule()) {
+            diagnostic.amend(Edit::replacement(
+                "...".to_string(),
+                value.location,
+                value.end_location.unwrap(),
+            ));
+        }
+
+        checker.diagnostics.push(diagnostic);
+    }
+}
diff --git a/...ruff/src/rules/flake8_pyi/snapshots/ruff__rules__flake8_pyi__tests__PYI015_PYI015.py.snap b/...ruff/src/rules/flake8_pyi/snapshots/ruff__rules__flake8_pyi__tests__PYI015_PYI015.py.snap
@@ -0,0 +1,6 @@
+---
+source: crates/ruff/src/rules/flake8_pyi/mod.rs
+expression: diagnostics
+---
+[]
+