Minor tweaks
This commit is contained in:
@@ -101,63 +101,6 @@ impl Violation for NonSelfReturnType {
|
||||
}
|
||||
}
|
||||
|
||||
fn import_self_edit(checker: &Checker, range: TextRange) -> Option<Edit> {
|
||||
let target_version = checker.settings.target_version.as_tuple();
|
||||
let source_module = if checker.source_type.is_stub() || target_version < (3, 11) {
|
||||
"typing_extensions"
|
||||
} else {
|
||||
"typing"
|
||||
};
|
||||
|
||||
let (importer, semantic) = (checker.importer(), checker.semantic());
|
||||
let request = ImportRequest::import_from(source_module, "Self");
|
||||
|
||||
let Ok((edit, ..)) = importer.get_or_import_symbol(&request, range.start(), semantic) else {
|
||||
return None;
|
||||
};
|
||||
|
||||
Some(edit)
|
||||
}
|
||||
|
||||
fn replace_with_self_fix(checker: &mut Checker, range: TextRange) -> Option<Fix> {
|
||||
let import_self = import_self_edit(checker, range)?;
|
||||
let replace_with_self = Edit::range_replacement("Self".to_string(), range);
|
||||
|
||||
Some(Fix::unsafe_edits(import_self, [replace_with_self]))
|
||||
}
|
||||
|
||||
fn add_diagnostic(
|
||||
checker: &mut Checker,
|
||||
stmt: &Stmt,
|
||||
returns: &Expr,
|
||||
class_name: String,
|
||||
method_name: String,
|
||||
) {
|
||||
let kind = NonSelfReturnType {
|
||||
class_name,
|
||||
method_name,
|
||||
};
|
||||
let mut diagnostic = Diagnostic::new(kind, stmt.identifier());
|
||||
|
||||
if let Some(fix) = replace_with_self_fix(checker, returns.range()) {
|
||||
diagnostic.set_fix(fix);
|
||||
}
|
||||
|
||||
checker.diagnostics.push(diagnostic);
|
||||
}
|
||||
|
||||
macro_rules! add_diagnostic {
|
||||
($checker:ident, $stmt:ident, $returns:ident, $class_def:ident, $method_name:ident) => {
|
||||
add_diagnostic(
|
||||
$checker,
|
||||
$stmt,
|
||||
$returns,
|
||||
$class_def.name.to_string(),
|
||||
$method_name.to_string(),
|
||||
);
|
||||
};
|
||||
}
|
||||
|
||||
/// PYI034
|
||||
pub(crate) fn non_self_return_type(
|
||||
checker: &mut Checker,
|
||||
@@ -197,7 +140,7 @@ pub(crate) fn non_self_return_type(
|
||||
&& is_name(returns, &class_def.name)
|
||||
&& !is_final(&class_def.decorator_list, semantic)
|
||||
{
|
||||
add_diagnostic!(checker, stmt, returns, class_def, name);
|
||||
add_diagnostic(checker, stmt, returns, class_def, name);
|
||||
}
|
||||
return;
|
||||
}
|
||||
@@ -205,7 +148,7 @@ pub(crate) fn non_self_return_type(
|
||||
// In-place methods that are expected to return `Self`.
|
||||
if is_inplace_bin_op(name) {
|
||||
if !is_self(returns, checker) {
|
||||
add_diagnostic!(checker, stmt, returns, class_def, name);
|
||||
add_diagnostic(checker, stmt, returns, class_def, name);
|
||||
}
|
||||
return;
|
||||
}
|
||||
@@ -213,7 +156,7 @@ pub(crate) fn non_self_return_type(
|
||||
if is_name(returns, &class_def.name) {
|
||||
if matches!(name, "__enter__" | "__new__") && !is_final(&class_def.decorator_list, semantic)
|
||||
{
|
||||
add_diagnostic!(checker, stmt, returns, class_def, name);
|
||||
add_diagnostic(checker, stmt, returns, class_def, name);
|
||||
}
|
||||
return;
|
||||
}
|
||||
@@ -223,20 +166,68 @@ pub(crate) fn non_self_return_type(
|
||||
if is_iterable_or_iterator(returns, semantic)
|
||||
&& subclasses_iterator(class_def, semantic)
|
||||
{
|
||||
add_diagnostic!(checker, stmt, returns, class_def, name);
|
||||
add_diagnostic(checker, stmt, returns, class_def, name);
|
||||
}
|
||||
}
|
||||
"__aiter__" => {
|
||||
if is_async_iterable_or_iterator(returns, semantic)
|
||||
&& subclasses_async_iterator(class_def, semantic)
|
||||
{
|
||||
add_diagnostic!(checker, stmt, returns, class_def, name);
|
||||
add_diagnostic(checker, stmt, returns, class_def, name);
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
/// Add a diagnostic for the given method.
|
||||
fn add_diagnostic(
|
||||
checker: &mut Checker,
|
||||
stmt: &Stmt,
|
||||
returns: &Expr,
|
||||
class_def: &ast::StmtClassDef,
|
||||
method_name: &str,
|
||||
) {
|
||||
/// Return an [`Edit`] that imports `typing.Self` from `typing` or `typing_extensions`.
|
||||
fn import_self(checker: &Checker, range: TextRange) -> Option<Edit> {
|
||||
let target_version = checker.settings.target_version.as_tuple();
|
||||
let source_module = if checker.source_type.is_stub() || target_version >= (3, 11) {
|
||||
"typing"
|
||||
} else {
|
||||
"typing_extensions"
|
||||
};
|
||||
|
||||
let (importer, semantic) = (checker.importer(), checker.semantic());
|
||||
let request = ImportRequest::import_from(source_module, "Self");
|
||||
|
||||
let Ok((edit, ..)) = importer.get_or_import_symbol(&request, range.start(), semantic)
|
||||
else {
|
||||
return None;
|
||||
};
|
||||
|
||||
Some(edit)
|
||||
}
|
||||
|
||||
/// Generate a [`Fix`] that replaces the return type with `Self`.
|
||||
fn replace_with_self(checker: &mut Checker, range: TextRange) -> Option<Fix> {
|
||||
let import_self = import_self(checker, range)?;
|
||||
let replace_with_self = Edit::range_replacement("Self".to_string(), range);
|
||||
Some(Fix::unsafe_edits(import_self, [replace_with_self]))
|
||||
}
|
||||
|
||||
let mut diagnostic = Diagnostic::new(
|
||||
NonSelfReturnType {
|
||||
class_name: class_def.name.to_string(),
|
||||
method_name: method_name.to_string(),
|
||||
},
|
||||
stmt.identifier(),
|
||||
);
|
||||
if let Some(fix) = replace_with_self(checker, returns.range()) {
|
||||
diagnostic.set_fix(fix);
|
||||
}
|
||||
checker.diagnostics.push(diagnostic);
|
||||
}
|
||||
|
||||
/// Returns `true` if the method is an in-place binary operator.
|
||||
fn is_inplace_bin_op(name: &str) -> bool {
|
||||
matches!(
|
||||
|
||||
Reference in New Issue
Block a user