Skip to content

Commit d01006e

Browse files
charliermarshcharliermarsh
committed
Tweaks
1 parent 8a4fe21 commit d01006e

File tree

2 files changed

+87
-78
lines changed

2 files changed

+87
-78
lines changed

‎crates/ruff/src/checkers/ast/analyze/statement.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -163,8 +163,8 @@ pub(crate) fn statement(stmt: &Stmt, checker: &mut Checker) {
163163
name,
164164
decorator_list,
165165
returns.as_ref().map(AsRef::as_ref),
166-
args,
167-
type_params,
166+
parameters,
167+
type_params.as_ref(),
168168
);
169169
}
170170
if checker.enabled(Rule::StrOrReprDefinedInStub) {

‎crates/ruff/src/rules/flake8_pyi/rules/custom_typevar_return.rs

Lines changed: 85 additions & 76 deletions
Original file line numberDiff line numberDiff line change
@@ -1,22 +1,27 @@
1-
use crate::checkers::ast::Checker;
2-
use crate::settings::types::PythonVersion;
31
use ruff_diagnostics::{Diagnostic, Violation};
42
use ruff_macros::{derive_message_formats, violation};
53
use ruff_python_ast as ast;
6-
use ruff_python_ast::{ArgWithDefault, Arguments, Decorator, Expr, Ranged, TypeParam};
4+
use ruff_python_ast::helpers::map_subscript;
5+
use ruff_python_ast::{
6+
Decorator, Expr, ParameterWithDefault, Parameters, Ranged, TypeParam, TypeParams,
7+
};
78
use ruff_python_semantic::analyze::visibility::{
8-
is_abstract, is_classmethod, is_overload, is_staticmethod,
9+
is_abstract, is_classmethod, is_new, is_overload, is_staticmethod,
910
};
10-
use ruff_python_semantic::ScopeKind;
11+
12+
use crate::checkers::ast::Checker;
1113

1214
/// ## What it does
13-
/// Checks if certain methods define a custom `TypeVar`s for their return annotation instead of
14-
/// using `typing_extensions.Self`. This check currently applies for instance methods that return
15-
/// `self`, class methods that return an instance of `cls`, and `__new__` methods.
15+
/// Checks for methods that define a custom `TypeVar` for their return type
16+
/// annotation instead of using `typing_extensions.Self`.
1617
///
1718
/// ## Why is this bad?
18-
/// If certain methods are annotated with a custom `TypeVar` type, and the class is subclassed,
19-
/// type checkers will not be able to infer the correct return type.
19+
/// If certain methods are annotated with a custom `TypeVar` type, and the
20+
/// class is subclassed, type checkers will not be able to infer the correct
21+
/// return type.
22+
///
23+
/// This check currently applies to instance methods that return `self`, class
24+
/// methods that return an instance of `cls`, and `__new__` methods.
2025
///
2126
/// ## Example
2227
/// ```python
@@ -58,7 +63,7 @@ impl Violation for CustomTypeVarReturnType {
5863
fn message(&self) -> String {
5964
let CustomTypeVarReturnType { method_name } = self;
6065
format!(
61-
"Methods like {method_name} should return `typing.Self` instead of a custom TypeVar"
66+
"Methods like `{method_name}` should return `typing.Self` instead of a custom `TypeVar`"
6267
)
6368
}
6469
}
@@ -69,13 +74,9 @@ pub(crate) fn custom_typevar_return_type(
6974
name: &str,
7075
decorator_list: &[Decorator],
7176
returns: Option<&Expr>,
72-
args: &Arguments,
73-
type_params: &[TypeParam],
77+
args: &Parameters,
78+
type_params: Option<&TypeParams>,
7479
) {
75-
let ScopeKind::Class(_) = checker.semantic().scope().kind else {
76-
return;
77-
};
78-
7980
if args.args.is_empty() && args.posonlyargs.is_empty() {
8081
return;
8182
}
@@ -84,120 +85,128 @@ pub(crate) fn custom_typevar_return_type(
8485
return;
8586
};
8687

87-
let return_annotation = if let Expr::Subscript(ast::ExprSubscript { value, .. }) = returns {
88-
// Ex) `Type[T]`
89-
value
90-
} else {
91-
// Ex) `Type`, `typing.Type`
92-
returns
88+
if !checker.semantic().scope().kind.is_class() {
89+
return;
9390
};
9491

95-
// Skip any abstract, static and overloaded methods.
92+
// Skip any abstract, static, and overloaded methods.
9693
if is_abstract(decorator_list, checker.semantic())
9794
|| is_overload(decorator_list, checker.semantic())
9895
|| is_staticmethod(decorator_list, checker.semantic())
9996
{
10097
return;
10198
}
10299

103-
let is_violation: bool =
104-
if is_classmethod(decorator_list, checker.semantic()) || name == "__new__" {
105-
check_class_method_for_bad_typevars(checker, args, return_annotation, type_params)
100+
let returns = map_subscript(returns);
101+
102+
let uses_custom_var: bool =
103+
if is_classmethod(decorator_list, checker.semantic()) || is_new(name) {
104+
class_method(args, returns, type_params)
106105
} else {
107106
// If not static, or a class method or __new__ we know it is an instance method
108-
check_instance_method_for_bad_typevars(checker, args, return_annotation, type_params)
107+
instance_method(args, returns, type_params)
109108
};
110109

111-
if is_violation {
110+
if uses_custom_var {
112111
checker.diagnostics.push(Diagnostic::new(
113112
CustomTypeVarReturnType {
114113
method_name: name.to_string(),
115114
},
116-
return_annotation.range(),
115+
returns.range(),
117116
));
118117
}
119118
}
120119

121-
fn check_class_method_for_bad_typevars(
122-
checker: &Checker,
123-
args: &Arguments,
120+
/// Returns `true` if the class method is annotated with a custom `TypeVar` that is likely
121+
/// private.
122+
fn class_method(
123+
args: &Parameters,
124124
return_annotation: &Expr,
125-
type_params: &[TypeParam],
125+
type_params: Option<&TypeParams>,
126126
) -> bool {
127-
let ArgWithDefault { def, .. } = &args.args[0];
127+
let ParameterWithDefault { parameter, .. } = &args.args[0];
128128

129-
let Some(annotation) = &def.annotation else {
130-
return false
129+
let Some(annotation) = &parameter.annotation else {
130+
return false;
131131
};
132132

133-
let Expr::Subscript(ast::ExprSubscript{slice, value, ..}) = annotation.as_ref() else {
134-
return false
133+
let Expr::Subscript(ast::ExprSubscript { slice, value, .. }) = annotation.as_ref() else {
134+
return false;
135135
};
136136

137-
let Expr::Name(ast::ExprName{ id: id_value, .. }) = value.as_ref() else {
138-
return false
137+
let Expr::Name(value) = value.as_ref() else {
138+
return false;
139139
};
140140

141-
// Don't error if the first argument is annotated with typing.Type[T]
141+
// Don't error if the first argument is annotated with typing.Type[T].
142142
// These are edge cases, and it's hard to give good error messages for them.
143-
if id_value != "type" {
143+
if value.id != "type" {
144144
return false;
145145
};
146146

147-
let Expr::Name(ast::ExprName { id: id_slice, .. }) = slice.as_ref() else {
148-
return false
147+
let Expr::Name(slice) = slice.as_ref() else {
148+
return false;
149149
};
150150

151-
let Expr::Name(ast::ExprName { id: return_type, .. }) = return_annotation else {
152-
return false
151+
let Expr::Name(return_annotation) = return_annotation else {
152+
return false;
153153
};
154154

155-
return_type == id_slice && is_likely_private_typevar(checker, id_slice, type_params)
155+
if slice.id != return_annotation.id {
156+
return false;
157+
}
158+
159+
is_likely_private_typevar(&slice.id, type_params)
156160
}
157161

158-
fn check_instance_method_for_bad_typevars(
159-
checker: &Checker,
160-
args: &Arguments,
162+
/// Returns `true` if the instance method is annotated with a custom `TypeVar` that is likely
163+
/// private.
164+
fn instance_method(
165+
args: &Parameters,
161166
return_annotation: &Expr,
162-
type_params: &[TypeParam],
167+
type_params: Option<&TypeParams>,
163168
) -> bool {
164-
let ArgWithDefault { def, .. } = &args.args[0];
169+
let ParameterWithDefault { parameter, .. } = &args.args[0];
165170

166-
let Some(annotation) = &def.annotation else {
167-
return false
171+
let Some(annotation) = &parameter.annotation else {
172+
return false;
168173
};
169174

170-
let Expr::Name(ast::ExprName{id: first_arg_type,..}) = annotation.as_ref() else {
171-
return false
175+
let Expr::Name(ast::ExprName {
176+
id: first_arg_type, ..
177+
}) = annotation.as_ref()
178+
else {
179+
return false;
172180
};
173181

174-
let Expr::Name(ast::ExprName { id: return_type, .. }) = return_annotation else {
175-
return false
182+
let Expr::Name(ast::ExprName {
183+
id: return_type, ..
184+
}) = return_annotation
185+
else {
186+
return false;
176187
};
177188

178189
if first_arg_type != return_type {
179190
return false;
180191
}
181192

182-
is_likely_private_typevar(checker, first_arg_type, type_params)
193+
is_likely_private_typevar(first_arg_type, type_params)
183194
}
184195

185-
fn is_likely_private_typevar(
186-
checker: &Checker,
187-
tvar_name: &str,
188-
type_params: &[TypeParam],
189-
) -> bool {
190-
if tvar_name.starts_with('_') {
196+
/// Returns `true` if the type variable is likely private.
197+
fn is_likely_private_typevar(type_var_name: &str, type_params: Option<&TypeParams>) -> bool {
198+
// Ex) `_T`
199+
if type_var_name.starts_with('_') {
191200
return true;
192201
}
193-
if checker.settings.target_version < PythonVersion::Py312 {
194-
return false;
195-
}
196-
197-
for param in type_params {
198-
if let TypeParam::TypeVar(ast::TypeParamTypeVar { name, .. }) = param {
199-
return name == tvar_name;
200-
}
201-
}
202-
false
202+
// Ex) `class Foo[T]: ...`
203+
type_params.is_some_and(|type_params| {
204+
type_params.iter().any(|type_param| {
205+
if let TypeParam::TypeVar(ast::TypeParamTypeVar { name, .. }) = type_param {
206+
name == type_var_name
207+
} else {
208+
false
209+
}
210+
})
211+
})
203212
}

0 commit comments

Comments
 (0)