Minor tweaks

This commit is contained in:
Charlie Marsh
2024-11-08 21:06:01 -05:00
parent 33945a5517
commit 947e1a4e4d

View File

@@ -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!(