Avoid panic with positional-only arguments in PYI019

astral-sh · Aug 4, 2023 · 08a7d27 · 08a7d27
1 parent 9bb2128
commit 08a7d27
Show file tree

Hide file tree

Showing 5 changed files with 55 additions and 41 deletions.
diff --git a/crates/ruff/resources/test/fixtures/flake8_pyi/PYI019.py b/crates/ruff/resources/test/fixtures/flake8_pyi/PYI019.py
@@ -14,6 +14,10 @@ def bad_instance_method(self: _S, arg: bytes) -> _S: ...  # PYI019
     def bad_class_method(cls: type[_S], arg: int) -> _S: ...  # PYI019
 
 
+    @classmethod
+    def bad_posonly_class_method(cls: type[_S], /) -> _S: ...  # PYI019
+
+
     @classmethod
     def excluded_edge_case(cls: Type[_S], arg: int) -> _S: ...  # Ok
 

diff --git a/crates/ruff/resources/test/fixtures/flake8_pyi/PYI019.pyi b/crates/ruff/resources/test/fixtures/flake8_pyi/PYI019.pyi
@@ -14,6 +14,10 @@ class BadClass:
     def bad_class_method(cls: type[_S], arg: int) -> _S: ...  # PYI019
 
 
+    @classmethod
+    def bad_posonly_class_method(cls: type[_S], /) -> _S: ...  # PYI019
+
+
     @classmethod
     def excluded_edge_case(cls: Type[_S], arg: int) -> _S: ...  # Ok
 

diff --git a/crates/ruff/src/rules/flake8_pyi/rules/custom_type_var_return_type.rs b/crates/ruff/src/rules/flake8_pyi/rules/custom_type_var_return_type.rs
@@ -2,9 +2,7 @@ use ruff_diagnostics::{Diagnostic, Violation};
 use ruff_macros::{derive_message_formats, violation};
 use ruff_python_ast as ast;
 use ruff_python_ast::helpers::map_subscript;
-use ruff_python_ast::{
-    Decorator, Expr, ParameterWithDefault, Parameters, Ranged, TypeParam, TypeParams,
-};
+use ruff_python_ast::{Decorator, Expr, Parameters, Ranged, TypeParam, TypeParams};
 use ruff_python_semantic::analyze::visibility::{
     is_abstract, is_classmethod, is_new, is_overload, is_staticmethod,
 };
@@ -77,11 +75,19 @@ pub(crate) fn custom_type_var_return_type(
     args: &Parameters,
     type_params: Option<&TypeParams>,
 ) {
-    if args.args.is_empty() && args.posonlyargs.is_empty() {
+    // Given, e.g., `def foo(self: _S, arg: bytes) -> _T`, extract `_T`.
+    let Some(return_annotation) = returns else {
         return;
-    }
+    };
 
-    let Some(returns) = returns else {
+    // Given, e.g., `def foo(self: _S, arg: bytes)`, extract `_S`.
+    let Some(self_or_cls_annotation) = args
+        .posonlyargs
+        .iter()
+        .chain(args.args.iter())
+        .next()
+        .and_then(|parameter_with_default| parameter_with_default.parameter.annotation.as_ref())
+    else {
         return;
     };
 
@@ -97,40 +103,32 @@ pub(crate) fn custom_type_var_return_type(
         return;
     }
 
-    let returns = map_subscript(returns);
-
     let uses_custom_var: bool =
         if is_classmethod(decorator_list, checker.semantic()) || is_new(name) {
-            class_method(args, returns, type_params)
+            class_method(self_or_cls_annotation, return_annotation, type_params)
         } else {
             // If not static, or a class method or __new__ we know it is an instance method
-            instance_method(args, returns, type_params)
+            instance_method(self_or_cls_annotation, return_annotation, type_params)
         };
 
     if uses_custom_var {
         checker.diagnostics.push(Diagnostic::new(
             CustomTypeVarReturnType {
                 method_name: name.to_string(),
             },
-            returns.range(),
+            return_annotation.range(),
         ));
     }
 }
 
 /// Returns `true` if the class method is annotated with a custom `TypeVar` that is likely
 /// private.
 fn class_method(
-    args: &Parameters,
+    cls_annotation: &Expr,
     return_annotation: &Expr,
     type_params: Option<&TypeParams>,
 ) -> bool {
-    let ParameterWithDefault { parameter, .. } = &args.args[0];
-
-    let Some(annotation) = &parameter.annotation else {
-        return false;
-    };
-
-    let Expr::Subscript(ast::ExprSubscript { slice, value, .. }) = annotation.as_ref() else {
+    let Expr::Subscript(ast::ExprSubscript { slice, value, .. }) = cls_annotation else {
         return false;
     };
 
@@ -148,7 +146,7 @@ fn class_method(
         return false;
     };
 
-    let Expr::Name(return_annotation) = return_annotation else {
+    let Expr::Name(return_annotation) = map_subscript(return_annotation) else {
         return false;
     };
 
@@ -162,26 +160,20 @@ fn class_method(
 /// Returns `true` if the instance method is annotated with a custom `TypeVar` that is likely
 /// private.
 fn instance_method(
-    args: &Parameters,
+    self_annotation: &Expr,
     return_annotation: &Expr,
     type_params: Option<&TypeParams>,
 ) -> bool {
-    let ParameterWithDefault { parameter, .. } = &args.args[0];
-
-    let Some(annotation) = &parameter.annotation else {
-        return false;
-    };
-
     let Expr::Name(ast::ExprName {
         id: first_arg_type, ..
-    }) = annotation.as_ref()
+    }) = self_annotation
     else {
         return false;
     };
 
     let Expr::Name(ast::ExprName {
         id: return_type, ..
-    }) = return_annotation
+    }) = map_subscript(return_annotation)
     else {
         return false;
     };

diff --git a/...ruff/src/rules/flake8_pyi/snapshots/ruff__rules__flake8_pyi__tests__PYI019_PYI019.py.snap b/...ruff/src/rules/flake8_pyi/snapshots/ruff__rules__flake8_pyi__tests__PYI019_PYI019.py.snap
@@ -21,17 +21,24 @@ PYI019.py:14:54: PYI019 Methods like `bad_class_method` should return `typing.Se
    |                                                      ^^ PYI019
    |
 
-PYI019.py:35:63: PYI019 Methods like `__new__` should return `typing.Self` instead of a custom `TypeVar`
+PYI019.py:18:55: PYI019 Methods like `bad_posonly_class_method` should return `typing.Self` instead of a custom `TypeVar`
    |
-33 | # Python > 3.12
-34 | class PEP695BadDunderNew[T]:
-35 |   def __new__[S](cls: type[S], *args: Any, ** kwargs: Any) -> S: ...  # PYI019
+17 |     @classmethod
+18 |     def bad_posonly_class_method(cls: type[_S], /) -> _S: ...  # PYI019
+   |                                                       ^^ PYI019
+   |
+
+PYI019.py:39:63: PYI019 Methods like `__new__` should return `typing.Self` instead of a custom `TypeVar`
+   |
+37 | # Python > 3.12
+38 | class PEP695BadDunderNew[T]:
+39 |   def __new__[S](cls: type[S], *args: Any, ** kwargs: Any) -> S: ...  # PYI019
    |                                                               ^ PYI019
    |
 
-PYI019.py:38:46: PYI019 Methods like `generic_instance_method` should return `typing.Self` instead of a custom `TypeVar`
+PYI019.py:42:46: PYI019 Methods like `generic_instance_method` should return `typing.Self` instead of a custom `TypeVar`
    |
-38 |   def generic_instance_method[S](self: S) -> S: ...  # PYI019
+42 |   def generic_instance_method[S](self: S) -> S: ...  # PYI019
    |                                              ^ PYI019
    |
 

diff --git a/...uff/src/rules/flake8_pyi/snapshots/ruff__rules__flake8_pyi__tests__PYI019_PYI019.pyi.snap b/...uff/src/rules/flake8_pyi/snapshots/ruff__rules__flake8_pyi__tests__PYI019_PYI019.pyi.snap
@@ -21,17 +21,24 @@ PYI019.pyi:14:54: PYI019 Methods like `bad_class_method` should return `typing.S
    |                                                      ^^ PYI019
    |
 
-PYI019.pyi:35:63: PYI019 Methods like `__new__` should return `typing.Self` instead of a custom `TypeVar`
+PYI019.pyi:18:55: PYI019 Methods like `bad_posonly_class_method` should return `typing.Self` instead of a custom `TypeVar`
    |
-33 | # Python > 3.12
-34 | class PEP695BadDunderNew[T]:
-35 |   def __new__[S](cls: type[S], *args: Any, ** kwargs: Any) -> S: ...  # PYI019
+17 |     @classmethod
+18 |     def bad_posonly_class_method(cls: type[_S], /) -> _S: ...  # PYI019
+   |                                                       ^^ PYI019
+   |
+
+PYI019.pyi:39:63: PYI019 Methods like `__new__` should return `typing.Self` instead of a custom `TypeVar`
+   |
+37 | # Python > 3.12
+38 | class PEP695BadDunderNew[T]:
+39 |   def __new__[S](cls: type[S], *args: Any, ** kwargs: Any) -> S: ...  # PYI019
    |                                                               ^ PYI019
    |
 
-PYI019.pyi:38:46: PYI019 Methods like `generic_instance_method` should return `typing.Self` instead of a custom `TypeVar`
+PYI019.pyi:42:46: PYI019 Methods like `generic_instance_method` should return `typing.Self` instead of a custom `TypeVar`
    |
-38 |   def generic_instance_method[S](self: S) -> S: ...  # PYI019
+42 |   def generic_instance_method[S](self: S) -> S: ...  # PYI019
    |                                              ^ PYI019
    |