From 08a7d27afcd6fd398605cbf03c1bcbd649ee2a64 Mon Sep 17 00:00:00 2001 From: Charlie Marsh Date: Fri, 4 Aug 2023 14:17:40 -0400 Subject: [PATCH] Avoid panic with positional-only arguments in PYI019 --- .../test/fixtures/flake8_pyi/PYI019.py | 4 ++ .../test/fixtures/flake8_pyi/PYI019.pyi | 4 ++ .../rules/custom_type_var_return_type.rs | 50 ++++++++----------- ...__flake8_pyi__tests__PYI019_PYI019.py.snap | 19 ++++--- ..._flake8_pyi__tests__PYI019_PYI019.pyi.snap | 19 ++++--- 5 files changed, 55 insertions(+), 41 deletions(-) diff --git a/crates/ruff/resources/test/fixtures/flake8_pyi/PYI019.py b/crates/ruff/resources/test/fixtures/flake8_pyi/PYI019.py index 37dd8dfb574c3..23af128a3848d 100644 --- 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 index 37dd8dfb574c3..23af128a3848d 100644 --- 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 index 4ba24634a64ec..6dd0728e77c1b 100644 --- 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,14 +103,12 @@ 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 { @@ -112,7 +116,7 @@ pub(crate) fn custom_type_var_return_type( CustomTypeVarReturnType { method_name: name.to_string(), }, - returns.range(), + return_annotation.range(), )); } } @@ -120,17 +124,11 @@ pub(crate) fn custom_type_var_return_type( /// 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) = ¶meter.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) = ¶meter.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/crates/ruff/src/rules/flake8_pyi/snapshots/ruff__rules__flake8_pyi__tests__PYI019_PYI019.py.snap b/crates/ruff/src/rules/flake8_pyi/snapshots/ruff__rules__flake8_pyi__tests__PYI019_PYI019.py.snap index d3a7e225a1dba..7ead8194b969f 100644 --- a/crates/ruff/src/rules/flake8_pyi/snapshots/ruff__rules__flake8_pyi__tests__PYI019_PYI019.py.snap +++ b/crates/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/crates/ruff/src/rules/flake8_pyi/snapshots/ruff__rules__flake8_pyi__tests__PYI019_PYI019.pyi.snap b/crates/ruff/src/rules/flake8_pyi/snapshots/ruff__rules__flake8_pyi__tests__PYI019_PYI019.pyi.snap index 121ee656304d4..03c5bbb5bed24 100644 --- a/crates/ruff/src/rules/flake8_pyi/snapshots/ruff__rules__flake8_pyi__tests__PYI019_PYI019.pyi.snap +++ b/crates/ruff/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 |