Categorize functions in pep8-naming (#624)

resources/test/fixtures/N804.py

@@ -1,19 +1,43 @@
+from abc import ABCMeta
+
+
 class Class:
-    @classmethod
+    def bad_method(this):
-    def bad_class_method(this):
+        pass
+
+    if False:
+
+        def extra_bad_method(this):
+            pass
+
+    def good_method(self):
         pass
 
     @classmethod
-    def good_class_method(cls):
+    def class_method(cls):
-        pass
-
-    def method(self):
         pass
 
     @staticmethod
     def static_method(x):
         return x
 
+    def __init__(self):
+        ...
+
+    def __new__(cls, *args, **kwargs):
+        ...
+
+    def __init_subclass__(self, default_name, **kwargs):
+        ...
+
+
+class MetaClass(ABCMeta):
+    def bad_method(self):
+        pass
+
+    def good_method(cls):
+        pass
+
 
 def func(x):
     return x

resources/test/fixtures/N805.py

@@ -1,11 +1,11 @@
-import random
+from abc import ABCMeta
 
 
 class Class:
     def bad_method(this):
         pass
 
-    if random.random(0, 2) == 0:
+    if False:
 
         def extra_bad_method(this):
             pass
@@ -21,6 +21,23 @@ class Class:
     def static_method(x):
         return x
 
+    def __init__(self):
+        ...
+
+    def __new__(cls, *args, **kwargs):
+        ...
+
+    def __init_subclass__(self, default_name, **kwargs):
+        ...
+
+
+class MetaClass(ABCMeta):
+    def bad_method(self):
+        pass
+
+    def good_method(cls):
+        pass
+
 
 def func(x):
     return x

src/ast/types.rs

@@ -1,6 +1,7 @@
 use std::collections::BTreeMap;
 use std::sync::atomic::{AtomicUsize, Ordering};
 
+use rustpython_ast::{Expr, Keyword};
 use rustpython_parser::ast::{Located, Location};
 
 fn id() -> usize {
@@ -30,9 +31,17 @@ pub struct FunctionScope {
     pub uses_locals: bool,
 }
 
+#[derive(Clone, Debug, Default)]
+pub struct ClassScope<'a> {
+    pub name: &'a str,
+    pub bases: &'a [Expr],
+    pub keywords: &'a [Keyword],
+    pub decorator_list: &'a [Expr],
+}
+
 #[derive(Clone, Debug)]
-pub enum ScopeKind {
+pub enum ScopeKind<'a> {
-    Class,
+    Class(ClassScope<'a>),
     Function(FunctionScope),
     Generator,
     Module,
@@ -40,15 +49,15 @@ pub enum ScopeKind {
 }
 
 #[derive(Clone, Debug)]
-pub struct Scope {
+pub struct Scope<'a> {
     pub id: usize,
-    pub kind: ScopeKind,
+    pub kind: ScopeKind<'a>,
     pub import_starred: bool,
     pub values: BTreeMap<String, Binding>,
 }
 
-impl Scope {
+impl<'a> Scope<'a> {
-    pub fn new(kind: ScopeKind) -> Self {
+    pub fn new(kind: ScopeKind<'a>) -> Self {
         Scope {
             id: id(),
             kind,

src/check_ast.rs

@@ -15,8 +15,8 @@ use crate::ast::helpers::{extract_handler_names, match_name_or_attr_from_module}
 use crate::ast::operations::extract_all_names;
 use crate::ast::relocate::relocate_expr;
 use crate::ast::types::{
-    Binding, BindingContext, BindingKind, CheckLocator, FunctionScope, ImportKind, Range, Scope,
+    Binding, BindingContext, BindingKind, CheckLocator, ClassScope, FunctionScope, ImportKind,
-    ScopeKind,
+    Range, Scope, ScopeKind,
 };
 use crate::ast::visitor::{walk_excepthandler, Visitor};
 use crate::ast::{helpers, operations, visitor};
@@ -66,7 +66,7 @@ pub struct Checker<'a> {
     // at various points in time.
     pub(crate) parents: Vec<&'a Stmt>,
     pub(crate) parent_stack: Vec<usize>,
-    scopes: Vec<Scope>,
+    scopes: Vec<Scope<'a>>,
     scope_stack: Vec<usize>,
     dead_scopes: Vec<usize>,
     deferred_string_annotations: Vec<(Range, &'a str)>,
@@ -274,6 +274,7 @@ where
                     if let Some(check) =
                         pep8_naming::checks::invalid_first_argument_name_for_class_method(
                             self.current_scope(),
+                            name,
                             decorator_list,
                             args,
                             &self.settings.pep8_naming,
@@ -286,6 +287,7 @@ where
                 if self.settings.enabled.contains(&CheckCode::N805) {
                     if let Some(check) = pep8_naming::checks::invalid_first_argument_name_for_method(
                         self.current_scope(),
+                        name,
                         decorator_list,
                         args,
                         &self.settings.pep8_naming,
@@ -296,7 +298,7 @@ where
 
                 if self.settings.enabled.contains(&CheckCode::N807) {
                     if let Some(check) =
-                        pep8_naming::checks::dunder_function_name(stmt, self.current_scope(), name)
+                        pep8_naming::checks::dunder_function_name(self.current_scope(), stmt, name)
                     {
                         self.checks.push(check);
                     }
@@ -360,7 +362,7 @@ where
                 if self.settings.enabled.contains(&CheckCode::F706) {
                     if let Some(scope_index) = self.scope_stack.last().cloned() {
                         match self.scopes[scope_index].kind {
-                            ScopeKind::Class | ScopeKind::Module => {
+                            ScopeKind::Class(_) | ScopeKind::Module => {
                                 self.checks.push(Check::new(
                                     CheckKind::ReturnOutsideFunction,
                                     self.locate_check(Range::from_located(stmt)),
@@ -377,7 +379,6 @@ where
                 keywords,
                 decorator_list,
                 body,
-                ..
             } => {
                 if self.settings.enabled.contains(&CheckCode::U004) {
                     pyupgrade::plugins::useless_object_inheritance(
@@ -427,7 +428,12 @@ where
                 for expr in decorator_list {
                     self.visit_expr(expr)
                 }
-                self.push_scope(Scope::new(ScopeKind::Class))
+                self.push_scope(Scope::new(ScopeKind::Class(ClassScope {
+                    name,
+                    bases,
+                    keywords,
+                    decorator_list,
+                })))
             }
             StmtKind::Import { names } => {
                 if self.settings.enabled.contains(&CheckCode::E402) {
@@ -1256,7 +1262,7 @@ where
             ExprKind::Yield { .. } | ExprKind::YieldFrom { .. } | ExprKind::Await { .. } => {
                 let scope = self.current_scope();
                 if self.settings.enabled.contains(&CheckCode::F704) {
-                    if matches!(scope.kind, ScopeKind::Class | ScopeKind::Module) {
+                    if matches!(scope.kind, ScopeKind::Class(_) | ScopeKind::Module) {
                         self.checks.push(Check::new(
                             CheckKind::YieldOutsideFunction,
                             self.locate_check(Range::from_located(expr)),
@@ -1791,7 +1797,7 @@ impl<'a> Checker<'a> {
             .expect("Attempted to pop without scope.");
     }
 
-    fn push_scope(&mut self, scope: Scope) {
+    fn push_scope(&mut self, scope: Scope<'a>) {
         self.scope_stack.push(self.scopes.len());
         self.scopes.push(scope);
     }
@@ -1896,7 +1902,7 @@ impl<'a> Checker<'a> {
             let mut import_starred = false;
             for scope_index in self.scope_stack.iter().rev() {
                 let scope = &mut self.scopes[*scope_index];
-                if matches!(scope.kind, ScopeKind::Class) {
+                if matches!(scope.kind, ScopeKind::Class(_)) {
                     if id == "__class__" {
                         return;
                     } else if !first_iter && !in_generator {
@@ -2433,7 +2439,7 @@ impl<'a> Checker<'a> {
     }
 
     fn check_builtin_shadowing(&mut self, name: &str, location: Range, is_attribute: bool) {
-        if is_attribute && matches!(self.current_scope().kind, ScopeKind::Class) {
+        if is_attribute && matches!(self.current_scope().kind, ScopeKind::Class(_)) {
             if self.settings.enabled.contains(&CheckCode::A003) {
                 if let Some(check) = flake8_builtins::checks::builtin_shadowing(
                     name,

src/pep8_naming/checks.rs

@@ -1,8 +1,9 @@
-use itertools::Itertools;
 use rustpython_ast::{Arguments, Expr, ExprKind, Stmt};
 
 use crate::ast::types::{FunctionScope, Range, Scope, ScopeKind};
 use crate::checks::{Check, CheckKind};
+use crate::pep8_naming::helpers;
+use crate::pep8_naming::helpers::FunctionType;
 use crate::pep8_naming::settings::Settings;
 
 /// N801
@@ -53,21 +54,15 @@ pub fn invalid_argument_name(name: &str, location: Range) -> Option<Check> {
 /// N804
 pub fn invalid_first_argument_name_for_class_method(
     scope: &Scope,
+    name: &str,
     decorator_list: &[Expr],
     args: &Arguments,
     settings: &Settings,
 ) -> Option<Check> {
-    if !matches!(scope.kind, ScopeKind::Class) {
+    if matches!(
-        return None;
+        helpers::function_type(scope, name, decorator_list, settings),
-    }
+        FunctionType::ClassMethod
-
+    ) {
-    if decorator_list.iter().any(|decorator| {
-        if let ExprKind::Name { id, .. } = &decorator.node {
-            settings.classmethod_decorators.contains(id)
-        } else {
-            false
-        }
-    }) {
         if let Some(arg) = args.args.first() {
             if arg.node.arg != "cls" {
                 return Some(Check::new(
@@ -83,31 +78,22 @@ pub fn invalid_first_argument_name_for_class_method(
 /// N805
 pub fn invalid_first_argument_name_for_method(
     scope: &Scope,
+    name: &str,
     decorator_list: &[Expr],
     args: &Arguments,
     settings: &Settings,
 ) -> Option<Check> {
-    if !matches!(scope.kind, ScopeKind::Class) {
+    if matches!(
-        return None;
+        helpers::function_type(scope, name, decorator_list, settings),
-    }
+        FunctionType::Method
-
+    ) {
-    if decorator_list.iter().any(|decorator| {
+        if let Some(arg) = args.args.first() {
-        if let ExprKind::Name { id, .. } = &decorator.node {
+            if arg.node.arg != "self" {
-            settings.classmethod_decorators.contains(id)
+                return Some(Check::new(
-                || settings.staticmethod_decorators.contains(id)
+                    CheckKind::InvalidFirstArgumentNameForMethod,
-        } else {
+                    Range::from_located(arg),
-            false
+                ));
-        }
+            }
-    }) {
-        return None;
-    }
-
-    if let Some(arg) = args.args.first() {
-        if arg.node.arg != "self" {
-            return Some(Check::new(
-                CheckKind::InvalidFirstArgumentNameForMethod,
-                Range::from_located(arg),
-            ));
         }
     }
     None
@@ -128,18 +114,16 @@ pub fn non_lowercase_variable_in_function(scope: &Scope, expr: &Expr, name: &str
 }
 
 /// N807
-pub fn dunder_function_name(func_def: &Stmt, scope: &Scope, name: &str) -> Option<Check> {
+pub fn dunder_function_name(scope: &Scope, stmt: &Stmt, name: &str) -> Option<Check> {
-    if matches!(scope.kind, ScopeKind::Class) {
+    if matches!(scope.kind, ScopeKind::Class(_)) {
         return None;
     }
-
     if name.starts_with("__") && name.ends_with("__") {
         return Some(Check::new(
             CheckKind::DunderFunctionName,
-            Range::from_located(func_def),
+            Range::from_located(stmt),
         ));
     }
-
     None
 }
 
@@ -149,7 +133,7 @@ pub fn constant_imported_as_non_constant(
     name: &str,
     asname: &str,
 ) -> Option<Check> {
-    if is_upper(name) && !is_upper(asname) {
+    if helpers::is_upper(name) && !helpers::is_upper(asname) {
         return Some(Check::new(
             CheckKind::ConstantImportedAsNonConstant(name.to_string(), asname.to_string()),
             Range::from_located(import_from),
@@ -164,7 +148,7 @@ pub fn lowercase_imported_as_non_lowercase(
     name: &str,
     asname: &str,
 ) -> Option<Check> {
-    if !is_upper(name) && is_lower(name) && asname.to_lowercase() != asname {
+    if !helpers::is_upper(name) && helpers::is_lower(name) && asname.to_lowercase() != asname {
         return Some(Check::new(
             CheckKind::LowercaseImportedAsNonLowercase(name.to_string(), asname.to_string()),
             Range::from_located(import_from),
@@ -179,7 +163,7 @@ pub fn camelcase_imported_as_lowercase(
     name: &str,
     asname: &str,
 ) -> Option<Check> {
-    if is_camelcase(name) && is_lower(asname) {
+    if helpers::is_camelcase(name) && helpers::is_lower(asname) {
         return Some(Check::new(
             CheckKind::CamelcaseImportedAsLowercase(name.to_string(), asname.to_string()),
             Range::from_located(import_from),
@@ -194,7 +178,11 @@ pub fn camelcase_imported_as_constant(
     name: &str,
     asname: &str,
 ) -> Option<Check> {
-    if is_camelcase(name) && !is_lower(asname) && is_upper(asname) && !is_acronym(name, asname) {
+    if helpers::is_camelcase(name)
+        && !helpers::is_lower(asname)
+        && helpers::is_upper(asname)
+        && !helpers::is_acronym(name, asname)
+    {
         return Some(Check::new(
             CheckKind::CamelcaseImportedAsConstant(name.to_string(), asname.to_string()),
             Range::from_located(import_from),
@@ -205,10 +193,10 @@ pub fn camelcase_imported_as_constant(
 
 /// N815
 pub fn mixed_case_variable_in_class_scope(scope: &Scope, expr: &Expr, name: &str) -> Option<Check> {
-    if !matches!(scope.kind, ScopeKind::Class) {
+    if !matches!(scope.kind, ScopeKind::Class(_)) {
         return None;
     }
-    if is_mixed_case(name) {
+    if helpers::is_mixed_case(name) {
         return Some(Check::new(
             CheckKind::MixedCaseVariableInClassScope(name.to_string()),
             Range::from_located(expr),
@@ -226,7 +214,7 @@ pub fn mixed_case_variable_in_global_scope(
     if !matches!(scope.kind, ScopeKind::Module) {
         return None;
     }
-    if is_mixed_case(name) {
+    if helpers::is_mixed_case(name) {
         return Some(Check::new(
             CheckKind::MixedCaseVariableInGlobalScope(name.to_string()),
             Range::from_located(expr),
@@ -241,7 +229,11 @@ pub fn camelcase_imported_as_acronym(
     name: &str,
     asname: &str,
 ) -> Option<Check> {
-    if is_camelcase(name) && !is_lower(asname) && is_upper(asname) && is_acronym(name, asname) {
+    if helpers::is_camelcase(name)
+        && !helpers::is_lower(asname)
+        && helpers::is_upper(asname)
+        && helpers::is_acronym(name, asname)
+    {
         return Some(Check::new(
             CheckKind::CamelcaseImportedAsAcronym(name.to_string(), asname.to_string()),
             Range::from_located(import_from),
@@ -272,101 +264,3 @@ pub fn error_suffix_on_exception_name(
     }
     None
 }
-
-fn is_lower(s: &str) -> bool {
-    let mut cased = false;
-    for c in s.chars() {
-        if c.is_uppercase() {
-            return false;
-        } else if !cased && c.is_lowercase() {
-            cased = true;
-        }
-    }
-    cased
-}
-
-fn is_upper(s: &str) -> bool {
-    let mut cased = false;
-    for c in s.chars() {
-        if c.is_lowercase() {
-            return false;
-        } else if !cased && c.is_uppercase() {
-            cased = true;
-        }
-    }
-    cased
-}
-
-fn is_camelcase(name: &str) -> bool {
-    !is_lower(name) && !is_upper(name) && !name.contains('_')
-}
-
-fn is_mixed_case(name: &str) -> bool {
-    !is_lower(name)
-        && name
-            .strip_prefix('_')
-            .unwrap_or(name)
-            .chars()
-            .next()
-            .map_or_else(|| false, |c| c.is_lowercase())
-}
-
-fn is_acronym(name: &str, asname: &str) -> bool {
-    name.chars().filter(|c| c.is_uppercase()).join("") == asname
-}
-
-#[cfg(test)]
-mod tests {
-    use super::{is_acronym, is_camelcase, is_lower, is_mixed_case, is_upper};
-
-    #[test]
-    fn test_is_lower() -> () {
-        assert!(is_lower("abc"));
-        assert!(is_lower("a_b_c"));
-        assert!(is_lower("a2c"));
-        assert!(!is_lower("aBc"));
-        assert!(!is_lower("ABC"));
-        assert!(!is_lower(""));
-        assert!(!is_lower("_"));
-    }
-
-    #[test]
-    fn test_is_upper() -> () {
-        assert!(is_upper("ABC"));
-        assert!(is_upper("A_B_C"));
-        assert!(is_upper("A2C"));
-        assert!(!is_upper("aBc"));
-        assert!(!is_upper("abc"));
-        assert!(!is_upper(""));
-        assert!(!is_upper("_"));
-    }
-
-    #[test]
-    fn test_is_camelcase() -> () {
-        assert!(is_camelcase("Camel"));
-        assert!(is_camelcase("CamelCase"));
-        assert!(!is_camelcase("camel"));
-        assert!(!is_camelcase("camel_case"));
-        assert!(!is_camelcase("CAMEL"));
-        assert!(!is_camelcase("CAMEL_CASE"));
-    }
-
-    #[test]
-    fn test_is_mixed_case() -> () {
-        assert!(is_mixed_case("mixedCase"));
-        assert!(is_mixed_case("mixed_Case"));
-        assert!(is_mixed_case("_mixed_Case"));
-        assert!(!is_mixed_case("mixed_case"));
-        assert!(!is_mixed_case("MIXED_CASE"));
-        assert!(!is_mixed_case(""));
-        assert!(!is_mixed_case("_"));
-    }
-
-    #[test]
-    fn test_is_acronym() -> () {
-        assert!(is_acronym("AB", "AB"));
-        assert!(is_acronym("AbcDef", "AD"));
-        assert!(!is_acronym("AbcDef", "Ad"));
-        assert!(!is_acronym("AbcDef", "AB"));
-    }
-}

src/pep8_naming/helpers.rs

@@ -0,0 +1,160 @@
+use itertools::Itertools;
+use rustpython_ast::{Expr, ExprKind};
+
+use crate::ast::helpers::match_name_or_attr;
+use crate::ast::types::{Scope, ScopeKind};
+use crate::pep8_naming::settings::Settings;
+
+const CLASS_METHODS: [&str; 3] = ["__new__", "__init_subclass__", "__class_getitem__"];
+const METACLASS_BASES: [&str; 2] = ["type", "ABCMeta"];
+
+pub enum FunctionType {
+    Function,
+    Method,
+    ClassMethod,
+    StaticMethod,
+}
+
+/// Classify a function based on its scope, name, and decorators.
+pub fn function_type(
+    scope: &Scope,
+    name: &str,
+    decorator_list: &[Expr],
+    settings: &Settings,
+) -> FunctionType {
+    if let ScopeKind::Class(scope) = &scope.kind {
+        // Special-case class method, like `__new__`.
+        if CLASS_METHODS.contains(&name)
+            // The class itself extends a known metaclass, so all methods are class methods.
+            || scope.bases.iter().any(|expr| {
+                METACLASS_BASES
+                    .iter()
+                    .any(|target| match_name_or_attr(expr, target))
+            })
+            // The method is decorated with a class method decorator (like `@classmethod`).
+            || decorator_list.iter().any(|expr| {
+            if let ExprKind::Name { id, .. } = &expr.node {
+                settings.classmethod_decorators.contains(id)
+            } else {
+                false
+            }
+        }) {
+            FunctionType::ClassMethod
+        } else if decorator_list.iter().any(|expr| {
+            if let ExprKind::Name { id, .. } = &expr.node {
+                settings.staticmethod_decorators.contains(id)
+            } else {
+                false
+            }
+        }) {
+            // The method is decorated with a static method decorator (like
+            // `@staticmethod`).
+            FunctionType::StaticMethod
+        } else {
+            // It's an instance method.
+            FunctionType::Method
+        }
+    } else {
+        FunctionType::Function
+    }
+}
+
+pub fn is_lower(s: &str) -> bool {
+    let mut cased = false;
+    for c in s.chars() {
+        if c.is_uppercase() {
+            return false;
+        } else if !cased && c.is_lowercase() {
+            cased = true;
+        }
+    }
+    cased
+}
+
+pub fn is_upper(s: &str) -> bool {
+    let mut cased = false;
+    for c in s.chars() {
+        if c.is_lowercase() {
+            return false;
+        } else if !cased && c.is_uppercase() {
+            cased = true;
+        }
+    }
+    cased
+}
+
+pub fn is_camelcase(name: &str) -> bool {
+    !is_lower(name) && !is_upper(name) && !name.contains('_')
+}
+
+pub fn is_mixed_case(name: &str) -> bool {
+    !is_lower(name)
+        && name
+            .strip_prefix('_')
+            .unwrap_or(name)
+            .chars()
+            .next()
+            .map_or_else(|| false, |c| c.is_lowercase())
+}
+
+pub fn is_acronym(name: &str, asname: &str) -> bool {
+    name.chars().filter(|c| c.is_uppercase()).join("") == asname
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::pep8_naming::helpers::{
+        is_acronym, is_camelcase, is_lower, is_mixed_case, is_upper,
+    };
+
+    #[test]
+    fn test_is_lower() -> () {
+        assert!(is_lower("abc"));
+        assert!(is_lower("a_b_c"));
+        assert!(is_lower("a2c"));
+        assert!(!is_lower("aBc"));
+        assert!(!is_lower("ABC"));
+        assert!(!is_lower(""));
+        assert!(!is_lower("_"));
+    }
+
+    #[test]
+    fn test_is_upper() -> () {
+        assert!(is_upper("ABC"));
+        assert!(is_upper("A_B_C"));
+        assert!(is_upper("A2C"));
+        assert!(!is_upper("aBc"));
+        assert!(!is_upper("abc"));
+        assert!(!is_upper(""));
+        assert!(!is_upper("_"));
+    }
+
+    #[test]
+    fn test_is_camelcase() -> () {
+        assert!(is_camelcase("Camel"));
+        assert!(is_camelcase("CamelCase"));
+        assert!(!is_camelcase("camel"));
+        assert!(!is_camelcase("camel_case"));
+        assert!(!is_camelcase("CAMEL"));
+        assert!(!is_camelcase("CAMEL_CASE"));
+    }
+
+    #[test]
+    fn test_is_mixed_case() -> () {
+        assert!(is_mixed_case("mixedCase"));
+        assert!(is_mixed_case("mixed_Case"));
+        assert!(is_mixed_case("_mixed_Case"));
+        assert!(!is_mixed_case("mixed_case"));
+        assert!(!is_mixed_case("MIXED_CASE"));
+        assert!(!is_mixed_case(""));
+        assert!(!is_mixed_case("_"));
+    }
+
+    #[test]
+    fn test_is_acronym() -> () {
+        assert!(is_acronym("AB", "AB"));
+        assert!(is_acronym("AbcDef", "AD"));
+        assert!(!is_acronym("AbcDef", "Ad"));
+        assert!(!is_acronym("AbcDef", "AB"));
+    }
+}

src/pep8_naming/mod.rs

@@ -1,2 +1,3 @@
 pub mod checks;
+mod helpers;
 pub mod settings;

src/snapshots/ruff__linter__tests__N804_N804.py.snap

@@ -4,10 +4,18 @@ expression: checks
 ---
 - kind: InvalidFirstArgumentNameForClassMethod
   location:
-    row: 3
+    row: 30
-    column: 25
+    column: 26
   end_location:
-    row: 3
+    row: 30
-    column: 29
+    column: 30
+  fix: ~
+- kind: InvalidFirstArgumentNameForClassMethod
+  location:
+    row: 35
+    column: 19
+  end_location:
+    row: 35
+    column: 23
   fix: ~