Treat type(Protocol) et al as metaclass base (#12770)

## Summary

Closes https://github.com/astral-sh/ruff/issues/12736.
This commit is contained in:
Charlie Marsh 2024-08-09 16:10:12 -04:00 committed by GitHub
parent 37b9bac403
commit 69e1c567d4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
11 changed files with 68 additions and 22 deletions

View file

@ -1,30 +1,40 @@
use rustc_hash::FxHashSet;
use crate::{BindingId, SemanticModel};
use ruff_python_ast as ast;
use ruff_python_ast::helpers::map_subscript;
use ruff_python_ast::name::QualifiedName;
use crate::{BindingId, SemanticModel};
use ruff_python_ast::Expr;
/// Return `true` if any base class matches a [`QualifiedName`] predicate.
pub fn any_qualified_name(
pub fn any_qualified_base_class(
class_def: &ast::StmtClassDef,
semantic: &SemanticModel,
func: &dyn Fn(QualifiedName) -> bool,
) -> bool {
any_base_class(class_def, semantic, &|expr| {
semantic
.resolve_qualified_name(map_subscript(expr))
.is_some_and(func)
})
}
/// Return `true` if any base class matches an [`Expr`] predicate.
pub fn any_base_class(
class_def: &ast::StmtClassDef,
semantic: &SemanticModel,
func: &dyn Fn(&Expr) -> bool,
) -> bool {
fn inner(
class_def: &ast::StmtClassDef,
semantic: &SemanticModel,
func: &dyn Fn(QualifiedName) -> bool,
func: &dyn Fn(&Expr) -> bool,
seen: &mut FxHashSet<BindingId>,
) -> bool {
class_def.bases().iter().any(|expr| {
// If the base class itself matches the pattern, then this does too.
// Ex) `class Foo(BaseModel): ...`
if semantic
.resolve_qualified_name(map_subscript(expr))
.is_some_and(func)
{
if func(expr) {
return true;
}
@ -100,7 +110,7 @@ pub fn any_super_class(
/// Return `true` if `class_def` is a class that has one or more enum classes in its mro
pub fn is_enumeration(class_def: &ast::StmtClassDef, semantic: &SemanticModel) -> bool {
any_qualified_name(class_def, semantic, &|qualified_name| {
any_qualified_base_class(class_def, semantic, &|qualified_name| {
matches!(
qualified_name.segments(),
[
@ -113,10 +123,26 @@ pub fn is_enumeration(class_def: &ast::StmtClassDef, semantic: &SemanticModel) -
/// Returns `true` if the given class is a metaclass.
pub fn is_metaclass(class_def: &ast::StmtClassDef, semantic: &SemanticModel) -> bool {
any_qualified_name(class_def, semantic, &|qualified_name| {
matches!(
qualified_name.segments(),
["" | "builtins", "type"] | ["abc", "ABCMeta"] | ["enum", "EnumMeta" | "EnumType"]
)
any_base_class(class_def, semantic, &|expr| match expr {
Expr::Call(ast::ExprCall {
func, arguments, ..
}) => {
// Ex) `class Foo(type(Protocol)): ...`
arguments.len() == 1 && semantic.match_builtin_expr(func.as_ref(), "type")
}
Expr::Subscript(ast::ExprSubscript { value, .. }) => {
// Ex) `class Foo(type[int]): ...`
semantic.match_builtin_expr(value.as_ref(), "type")
}
_ => semantic
.resolve_qualified_name(expr)
.is_some_and(|qualified_name| {
matches!(
qualified_name.segments(),
["" | "builtins", "type"]
| ["abc", "ABCMeta"]
| ["enum", "EnumMeta" | "EnumType"]
)
}),
})
}