Skip to content

Commit 5e73345

Browse files
Avoid panic with positional-only arguments in PYI019 (#6350)
## Summary Previously, failed on methods like: ```python @classmethod def bad_posonly_class_method(cls: type[_S], /) -> _S: ... # PYI019 ``` Since we check if there are any positional-only or non-positional arguments, but then do an unsafe access on `parameters.args`. Closes #6349. ## Test Plan `cargo test` (verified that `main` panics on the new fixtures)
1 parent b8fd693 commit 5e73345

File tree

5 files changed

+55
-41
lines changed

5 files changed

+55
-41
lines changed

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,10 @@ def bad_instance_method(self: _S, arg: bytes) -> _S: ... # PYI019
1414
def bad_class_method(cls: type[_S], arg: int) -> _S: ... # PYI019
1515

1616

17+
@classmethod
18+
def bad_posonly_class_method(cls: type[_S], /) -> _S: ... # PYI019
19+
20+
1721
@classmethod
1822
def excluded_edge_case(cls: Type[_S], arg: int) -> _S: ... # Ok
1923

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,10 @@ class BadClass:
1414
def bad_class_method(cls: type[_S], arg: int) -> _S: ... # PYI019
1515

1616

17+
@classmethod
18+
def bad_posonly_class_method(cls: type[_S], /) -> _S: ... # PYI019
19+
20+
1721
@classmethod
1822
def excluded_edge_case(cls: Type[_S], arg: int) -> _S: ... # Ok
1923

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

Lines changed: 21 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,7 @@ use ruff_diagnostics::{Diagnostic, Violation};
22
use ruff_macros::{derive_message_formats, violation};
33
use ruff_python_ast as ast;
44
use ruff_python_ast::helpers::map_subscript;
5-
use ruff_python_ast::{
6-
Decorator, Expr, ParameterWithDefault, Parameters, Ranged, TypeParam, TypeParams,
7-
};
5+
use ruff_python_ast::{Decorator, Expr, Parameters, Ranged, TypeParam, TypeParams};
86
use ruff_python_semantic::analyze::visibility::{
97
is_abstract, is_classmethod, is_new, is_overload, is_staticmethod,
108
};
@@ -77,11 +75,19 @@ pub(crate) fn custom_type_var_return_type(
7775
args: &Parameters,
7876
type_params: Option<&TypeParams>,
7977
) {
80-
if args.args.is_empty() && args.posonlyargs.is_empty() {
78+
// Given, e.g., `def foo(self: _S, arg: bytes) -> _T`, extract `_T`.
79+
let Some(return_annotation) = returns else {
8180
return;
82-
}
81+
};
8382

84-
let Some(returns) = returns else {
83+
// Given, e.g., `def foo(self: _S, arg: bytes)`, extract `_S`.
84+
let Some(self_or_cls_annotation) = args
85+
.posonlyargs
86+
.iter()
87+
.chain(args.args.iter())
88+
.next()
89+
.and_then(|parameter_with_default| parameter_with_default.parameter.annotation.as_ref())
90+
else {
8591
return;
8692
};
8793

@@ -97,40 +103,32 @@ pub(crate) fn custom_type_var_return_type(
97103
return;
98104
}
99105

100-
let returns = map_subscript(returns);
101-
102106
let uses_custom_var: bool =
103107
if is_classmethod(decorator_list, checker.semantic()) || is_new(name) {
104-
class_method(args, returns, type_params)
108+
class_method(self_or_cls_annotation, return_annotation, type_params)
105109
} else {
106110
// If not static, or a class method or __new__ we know it is an instance method
107-
instance_method(args, returns, type_params)
111+
instance_method(self_or_cls_annotation, return_annotation, type_params)
108112
};
109113

110114
if uses_custom_var {
111115
checker.diagnostics.push(Diagnostic::new(
112116
CustomTypeVarReturnType {
113117
method_name: name.to_string(),
114118
},
115-
returns.range(),
119+
return_annotation.range(),
116120
));
117121
}
118122
}
119123

120124
/// Returns `true` if the class method is annotated with a custom `TypeVar` that is likely
121125
/// private.
122126
fn class_method(
123-
args: &Parameters,
127+
cls_annotation: &Expr,
124128
return_annotation: &Expr,
125129
type_params: Option<&TypeParams>,
126130
) -> bool {
127-
let ParameterWithDefault { parameter, .. } = &args.args[0];
128-
129-
let Some(annotation) = &parameter.annotation else {
130-
return false;
131-
};
132-
133-
let Expr::Subscript(ast::ExprSubscript { slice, value, .. }) = annotation.as_ref() else {
131+
let Expr::Subscript(ast::ExprSubscript { slice, value, .. }) = cls_annotation else {
134132
return false;
135133
};
136134

@@ -148,7 +146,7 @@ fn class_method(
148146
return false;
149147
};
150148

151-
let Expr::Name(return_annotation) = return_annotation else {
149+
let Expr::Name(return_annotation) = map_subscript(return_annotation) else {
152150
return false;
153151
};
154152

@@ -162,26 +160,20 @@ fn class_method(
162160
/// Returns `true` if the instance method is annotated with a custom `TypeVar` that is likely
163161
/// private.
164162
fn instance_method(
165-
args: &Parameters,
163+
self_annotation: &Expr,
166164
return_annotation: &Expr,
167165
type_params: Option<&TypeParams>,
168166
) -> bool {
169-
let ParameterWithDefault { parameter, .. } = &args.args[0];
170-
171-
let Some(annotation) = &parameter.annotation else {
172-
return false;
173-
};
174-
175167
let Expr::Name(ast::ExprName {
176168
id: first_arg_type, ..
177-
}) = annotation.as_ref()
169+
}) = self_annotation
178170
else {
179171
return false;
180172
};
181173

182174
let Expr::Name(ast::ExprName {
183175
id: return_type, ..
184-
}) = return_annotation
176+
}) = map_subscript(return_annotation)
185177
else {
186178
return false;
187179
};

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

Lines changed: 13 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -21,17 +21,24 @@ PYI019.py:14:54: PYI019 Methods like `bad_class_method` should return `typing.Se
2121
| ^^ PYI019
2222
|
2323

24-
PYI019.py:35:63: PYI019 Methods like `__new__` should return `typing.Self` instead of a custom `TypeVar`
24+
PYI019.py:18:55: PYI019 Methods like `bad_posonly_class_method` should return `typing.Self` instead of a custom `TypeVar`
2525
|
26-
33 | # Python > 3.12
27-
34 | class PEP695BadDunderNew[T]:
28-
35 | def __new__[S](cls: type[S], *args: Any, ** kwargs: Any) -> S: ... # PYI019
26+
17 | @classmethod
27+
18 | def bad_posonly_class_method(cls: type[_S], /) -> _S: ... # PYI019
28+
| ^^ PYI019
29+
|
30+
31+
PYI019.py:39:63: PYI019 Methods like `__new__` should return `typing.Self` instead of a custom `TypeVar`
32+
|
33+
37 | # Python > 3.12
34+
38 | class PEP695BadDunderNew[T]:
35+
39 | def __new__[S](cls: type[S], *args: Any, ** kwargs: Any) -> S: ... # PYI019
2936
| ^ PYI019
3037
|
3138

32-
PYI019.py:38:46: PYI019 Methods like `generic_instance_method` should return `typing.Self` instead of a custom `TypeVar`
39+
PYI019.py:42:46: PYI019 Methods like `generic_instance_method` should return `typing.Self` instead of a custom `TypeVar`
3340
|
34-
38 | def generic_instance_method[S](self: S) -> S: ... # PYI019
41+
42 | def generic_instance_method[S](self: S) -> S: ... # PYI019
3542
| ^ PYI019
3643
|
3744

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

Lines changed: 13 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -21,17 +21,24 @@ PYI019.pyi:14:54: PYI019 Methods like `bad_class_method` should return `typing.S
2121
| ^^ PYI019
2222
|
2323

24-
PYI019.pyi:35:63: PYI019 Methods like `__new__` should return `typing.Self` instead of a custom `TypeVar`
24+
PYI019.pyi:18:55: PYI019 Methods like `bad_posonly_class_method` should return `typing.Self` instead of a custom `TypeVar`
2525
|
26-
33 | # Python > 3.12
27-
34 | class PEP695BadDunderNew[T]:
28-
35 | def __new__[S](cls: type[S], *args: Any, ** kwargs: Any) -> S: ... # PYI019
26+
17 | @classmethod
27+
18 | def bad_posonly_class_method(cls: type[_S], /) -> _S: ... # PYI019
28+
| ^^ PYI019
29+
|
30+
31+
PYI019.pyi:39:63: PYI019 Methods like `__new__` should return `typing.Self` instead of a custom `TypeVar`
32+
|
33+
37 | # Python > 3.12
34+
38 | class PEP695BadDunderNew[T]:
35+
39 | def __new__[S](cls: type[S], *args: Any, ** kwargs: Any) -> S: ... # PYI019
2936
| ^ PYI019
3037
|
3138

32-
PYI019.pyi:38:46: PYI019 Methods like `generic_instance_method` should return `typing.Self` instead of a custom `TypeVar`
39+
PYI019.pyi:42:46: PYI019 Methods like `generic_instance_method` should return `typing.Self` instead of a custom `TypeVar`
3340
|
34-
38 | def generic_instance_method[S](self: S) -> S: ... # PYI019
41+
42 | def generic_instance_method[S](self: S) -> S: ... # PYI019
3542
| ^ PYI019
3643
|
3744

0 commit comments

Comments
 (0)