@@ -2,9 +2,7 @@ use ruff_diagnostics::{Diagnostic, Violation};
2
2
use ruff_macros:: { derive_message_formats, violation} ;
3
3
use ruff_python_ast as ast;
4
4
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 } ;
8
6
use ruff_python_semantic:: analyze:: visibility:: {
9
7
is_abstract, is_classmethod, is_new, is_overload, is_staticmethod,
10
8
} ;
@@ -77,11 +75,19 @@ pub(crate) fn custom_type_var_return_type(
77
75
args : & Parameters ,
78
76
type_params : Option < & TypeParams > ,
79
77
) {
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 {
81
80
return ;
82
- }
81
+ } ;
83
82
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 {
85
91
return ;
86
92
} ;
87
93
@@ -97,40 +103,32 @@ pub(crate) fn custom_type_var_return_type(
97
103
return ;
98
104
}
99
105
100
- let returns = map_subscript ( returns) ;
101
-
102
106
let uses_custom_var: bool =
103
107
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)
105
109
} else {
106
110
// 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)
108
112
} ;
109
113
110
114
if uses_custom_var {
111
115
checker. diagnostics . push ( Diagnostic :: new (
112
116
CustomTypeVarReturnType {
113
117
method_name : name. to_string ( ) ,
114
118
} ,
115
- returns . range ( ) ,
119
+ return_annotation . range ( ) ,
116
120
) ) ;
117
121
}
118
122
}
119
123
120
124
/// Returns `true` if the class method is annotated with a custom `TypeVar` that is likely
121
125
/// private.
122
126
fn class_method (
123
- args : & Parameters ,
127
+ cls_annotation : & Expr ,
124
128
return_annotation : & Expr ,
125
129
type_params : Option < & TypeParams > ,
126
130
) -> 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 {
134
132
return false ;
135
133
} ;
136
134
@@ -148,7 +146,7 @@ fn class_method(
148
146
return false ;
149
147
} ;
150
148
151
- let Expr :: Name ( return_annotation) = return_annotation else {
149
+ let Expr :: Name ( return_annotation) = map_subscript ( return_annotation) else {
152
150
return false ;
153
151
} ;
154
152
@@ -162,26 +160,20 @@ fn class_method(
162
160
/// Returns `true` if the instance method is annotated with a custom `TypeVar` that is likely
163
161
/// private.
164
162
fn instance_method (
165
- args : & Parameters ,
163
+ self_annotation : & Expr ,
166
164
return_annotation : & Expr ,
167
165
type_params : Option < & TypeParams > ,
168
166
) -> bool {
169
- let ParameterWithDefault { parameter, .. } = & args. args [ 0 ] ;
170
-
171
- let Some ( annotation) = & parameter. annotation else {
172
- return false ;
173
- } ;
174
-
175
167
let Expr :: Name ( ast:: ExprName {
176
168
id : first_arg_type, ..
177
- } ) = annotation . as_ref ( )
169
+ } ) = self_annotation
178
170
else {
179
171
return false ;
180
172
} ;
181
173
182
174
let Expr :: Name ( ast:: ExprName {
183
175
id : return_type, ..
184
- } ) = return_annotation
176
+ } ) = map_subscript ( return_annotation)
185
177
else {
186
178
return false ;
187
179
} ;
0 commit comments