Treat type(Protocol) et al as metaclass base

astral-sh · Aug 9, 2024 · 8d67ca4 · 8d67ca4
1 parent bc5b9b8
commit 8d67ca4
Show file tree

Hide file tree

Showing 11 changed files with 120 additions and 55 deletions.
diff --git a/crates/ruff_linter/resources/test/fixtures/pep8_naming/N805.py b/crates/ruff_linter/resources/test/fixtures/pep8_naming/N805.py
@@ -123,3 +123,14 @@ def bad_method(this):
 class RenamingWithNFKC:
     def formula(household):
         hºusehold(1)
+
+
+from typing import Protocol
+
+
+class MyMeta(type):
+    def __subclasscheck__(cls, other): ...
+
+
+class MyProtocolMeta(type(Protocol)):
+    def __subclasscheck__(cls, other): ...
diff --git a/crates/ruff_linter/src/rules/flake8_django/rules/helpers.rs b/crates/ruff_linter/src/rules/flake8_django/rules/helpers.rs
@@ -1,24 +1,32 @@
+use ruff_python_ast::helpers::map_subscript;
 use ruff_python_ast::{self as ast, Expr};
-
 use ruff_python_semantic::{analyze, SemanticModel};
 
 /// Return `true` if a Python class appears to be a Django model, based on its base classes.
 pub(super) fn is_model(class_def: &ast::StmtClassDef, semantic: &SemanticModel) -> bool {
-    analyze::class::any_qualified_name(class_def, semantic, &|qualified_name| {
-        matches!(
-            qualified_name.segments(),
-            ["django", "db", "models", "Model"]
-        )
+    analyze::class::any_qualified_name(class_def, semantic, &|expr| {
+        semantic
+            .resolve_qualified_name(map_subscript(expr))
+            .is_some_and(|qualified_name| {
+                matches!(
+                    qualified_name.segments(),
+                    ["django", "db", "models", "Model"]
+                )
+            })
     })
 }
 
 /// Return `true` if a Python class appears to be a Django model form, based on its base classes.
 pub(super) fn is_model_form(class_def: &ast::StmtClassDef, semantic: &SemanticModel) -> bool {
-    analyze::class::any_qualified_name(class_def, semantic, &|qualified_name| {
-        matches!(
-            qualified_name.segments(),
-            ["django", "forms", "ModelForm"] | ["django", "forms", "models", "ModelForm"]
-        )
+    analyze::class::any_qualified_name(class_def, semantic, &|expr| {
+        semantic
+            .resolve_qualified_name(map_subscript(expr))
+            .is_some_and(|qualified_name| {
+                matches!(
+                    qualified_name.segments(),
+                    ["django", "forms", "ModelForm"] | ["django", "forms", "models", "ModelForm"]
+                )
+            })
     })
 }
 

diff --git a/crates/ruff_linter/src/rules/flake8_pyi/rules/non_self_return_type.rs b/crates/ruff_linter/src/rules/flake8_pyi/rules/non_self_return_type.rs
@@ -254,11 +254,15 @@ fn is_self(expr: &Expr, semantic: &SemanticModel) -> bool {
 
 /// Return `true` if the given class extends `collections.abc.Iterator`.
 fn subclasses_iterator(class_def: &ast::StmtClassDef, semantic: &SemanticModel) -> bool {
-    analyze::class::any_qualified_name(class_def, semantic, &|qualified_name| {
-        matches!(
-            qualified_name.segments(),
-            ["typing", "Iterator"] | ["collections", "abc", "Iterator"]
-        )
+    analyze::class::any_qualified_name(class_def, semantic, &|expr| {
+        semantic
+            .resolve_qualified_name(map_subscript(expr))
+            .is_some_and(|qualified_name| {
+                matches!(
+                    qualified_name.segments(),
+                    ["typing", "Iterator"] | ["collections", "abc", "Iterator"]
+                )
+            })
     })
 }
 
@@ -277,11 +281,15 @@ fn is_iterable_or_iterator(expr: &Expr, semantic: &SemanticModel) -> bool {
 
 /// Return `true` if the given class extends `collections.abc.AsyncIterator`.
 fn subclasses_async_iterator(class_def: &ast::StmtClassDef, semantic: &SemanticModel) -> bool {
-    analyze::class::any_qualified_name(class_def, semantic, &|qualified_name| {
-        matches!(
-            qualified_name.segments(),
-            ["typing", "AsyncIterator"] | ["collections", "abc", "AsyncIterator"]
-        )
+    analyze::class::any_qualified_name(class_def, semantic, &|expr| {
+        semantic
+            .resolve_qualified_name(map_subscript(expr))
+            .is_some_and(|qualified_name| {
+                matches!(
+                    qualified_name.segments(),
+                    ["typing", "AsyncIterator"] | ["collections", "abc", "AsyncIterator"]
+                )
+            })
     })
 }
 

diff --git a/crates/ruff_linter/src/rules/flake8_type_checking/helpers.rs b/crates/ruff_linter/src/rules/flake8_type_checking/helpers.rs
@@ -78,10 +78,14 @@ fn runtime_required_base_class(
     base_classes: &[String],
     semantic: &SemanticModel,
 ) -> bool {
-    analyze::class::any_qualified_name(class_def, semantic, &|qualified_name| {
-        base_classes
-            .iter()
-            .any(|base_class| QualifiedName::from_dotted_name(base_class) == qualified_name)
+    analyze::class::any_qualified_name(class_def, semantic, &|expr| {
+        semantic
+            .resolve_qualified_name(map_subscript(expr))
+            .is_some_and(|qualified_name| {
+                base_classes
+                    .iter()
+                    .any(|base_class| QualifiedName::from_dotted_name(base_class) == qualified_name)
+            })
     })
 }
 

diff --git a/crates/ruff_linter/src/rules/pep8_naming/helpers.rs b/crates/ruff_linter/src/rules/pep8_naming/helpers.rs
@@ -91,8 +91,8 @@ pub(super) fn is_typed_dict_class(class_def: &ast::StmtClassDef, semantic: &Sema
         return false;
     }
 
-    analyze::class::any_qualified_name(class_def, semantic, &|qualified_name| {
-        semantic.match_typing_qualified_name(&qualified_name, "TypedDict")
+    analyze::class::any_qualified_name(class_def, semantic, &|expr| {
+        semantic.match_typing_expr(expr, "TypedDict")
     })
 }
 

diff --git a/...src/rules/pep8_naming/snapshots/ruff_linter__rules__pep8_naming__tests__N805_N805.py.snap b/...src/rules/pep8_naming/snapshots/ruff_linter__rules__pep8_naming__tests__N805_N805.py.snap
@@ -286,3 +286,6 @@ N805.py:124:17: N805 [*] First argument of a method should be named `self`
 125     |-        hºusehold(1)
     124 |+    def formula(self):
     125 |+        self(1)
+126 126 | 
+127 127 | 
+128 128 | from typing import Protocol
diff --git a/...pep8_naming/snapshots/ruff_linter__rules__pep8_naming__tests__classmethod_decorators.snap b/...pep8_naming/snapshots/ruff_linter__rules__pep8_naming__tests__classmethod_decorators.snap
@@ -229,3 +229,6 @@ N805.py:124:17: N805 [*] First argument of a method should be named `self`
 125     |-        hºusehold(1)
     124 |+    def formula(self):
     125 |+        self(1)
+126 126 | 
+127 127 | 
+128 128 | from typing import Protocol
diff --git a/...ep8_naming/snapshots/ruff_linter__rules__pep8_naming__tests__staticmethod_decorators.snap b/...ep8_naming/snapshots/ruff_linter__rules__pep8_naming__tests__staticmethod_decorators.snap
@@ -267,3 +267,6 @@ N805.py:124:17: N805 [*] First argument of a method should be named `self`
 125     |-        hºusehold(1)
     124 |+    def formula(self):
     125 |+        self(1)
+126 126 | 
+127 127 | 
+128 128 | from typing import Protocol
diff --git a/crates/ruff_linter/src/rules/pydoclint/rules/check_docstring.rs b/crates/ruff_linter/src/rules/pydoclint/rules/check_docstring.rs
@@ -519,7 +519,7 @@ struct BodyEntries<'a> {
 struct BodyVisitor<'a> {
     returns: Vec<Entry>,
     yields: Vec<Entry>,
-    currently_suspended_exceptions: Option<&'a ast::Expr>,
+    currently_suspended_exceptions: Option<&'a Expr>,
     raised_exceptions: Vec<ExceptionEntry<'a>>,
     semantic: &'a SemanticModel<'a>,
 }

diff --git a/crates/ruff_linter/src/rules/ruff/rules/helpers.rs b/crates/ruff_linter/src/rules/ruff/rules/helpers.rs
@@ -74,13 +74,17 @@ pub(super) fn has_default_copy_semantics(
     class_def: &ast::StmtClassDef,
     semantic: &SemanticModel,
 ) -> bool {
-    analyze::class::any_qualified_name(class_def, semantic, &|qualified_name| {
-        matches!(
-            qualified_name.segments(),
-            ["pydantic", "BaseModel" | "BaseSettings" | "BaseConfig"]
-                | ["pydantic_settings", "BaseSettings"]
-                | ["msgspec", "Struct"]
-        )
+    analyze::class::any_qualified_name(class_def, semantic, &|expr| {
+        semantic
+            .resolve_qualified_name(map_subscript(expr))
+            .is_some_and(|qualified_name| {
+                matches!(
+                    qualified_name.segments(),
+                    ["pydantic", "BaseModel" | "BaseSettings" | "BaseConfig"]
+                        | ["pydantic_settings", "BaseSettings"]
+                        | ["msgspec", "Struct"]
+                )
+            })
     })
 }
 

diff --git a/crates/ruff_python_semantic/src/analyze/class.rs b/crates/ruff_python_semantic/src/analyze/class.rs
@@ -2,29 +2,26 @@ use rustc_hash::FxHashSet;
 
 use ruff_python_ast as ast;
 use ruff_python_ast::helpers::map_subscript;
-use ruff_python_ast::name::QualifiedName;
+use ruff_python_ast::Expr;
 
 use crate::{BindingId, SemanticModel};
 
 /// Return `true` if any base class matches a [`QualifiedName`] predicate.
 pub fn any_qualified_name(
     class_def: &ast::StmtClassDef,
     semantic: &SemanticModel,
-    func: &dyn Fn(QualifiedName) -> bool,
+    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,23 +97,47 @@ 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| {
-        matches!(
-            qualified_name.segments(),
-            [
-                "enum",
-                "Enum" | "Flag" | "IntEnum" | "IntFlag" | "StrEnum" | "ReprEnum" | "CheckEnum"
-            ]
-        )
+    any_qualified_name(class_def, semantic, &|expr| {
+        semantic
+            .resolve_qualified_name(map_subscript(expr))
+            .is_some_and(|qualified_name| {
+                matches!(
+                    qualified_name.segments(),
+                    [
+                        "enum",
+                        "Enum"
+                            | "Flag"
+                            | "IntEnum"
+                            | "IntFlag"
+                            | "StrEnum"
+                            | "ReprEnum"
+                            | "CheckEnum"
+                    ]
+                )
+            })
     })
 }
 
 /// 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_qualified_name(class_def, semantic, &|expr| match expr {
+        Expr::Call(ast::ExprCall { func, .. }) => {
+            // Ex) `class Foo(type(Protocol)): ...`
+            semantic
+                .resolve_qualified_name(func.as_ref())
+                .is_some_and(|qualified_name| {
+                    matches!(qualified_name.segments(), ["" | "builtins", "type"])
+                })
+        }
+        _ => semantic
+            .resolve_qualified_name(map_subscript(expr))
+            .is_some_and(|qualified_name| {
+                matches!(
+                    qualified_name.segments(),
+                    ["" | "builtins", "type"]
+                        | ["abc", "ABCMeta"]
+                        | ["enum", "EnumMeta" | "EnumType"]
+                )
+            }),
     })
 }