Refactor metaclass logic a bit, ensure lambdas are excluded

python/ql/src/Functions/MethodArgNames.qll

@@ -2,28 +2,19 @@
 
 import python
 import semmle.python.ApiGraphs
+import semmle.python.dataflow.new.internal.DataFlowDispatch
 import DataFlow
 
 /** Holds if `f` is a method of the class `c`. */
-private predicate methodOfClass(Function f, Class c) { f.getScope() = c }
+private predicate methodOfClass(Function f, Class c) {
+  exists(FunctionDef d | d.getDefinedFunction() = f and d.getScope() = c)
+}
 
 /** Holds if `c` is a metaclass. */
 private predicate isMetaclass(Class c) {
   c = API::builtin("type").getASubclass*().asSource().asExpr().(ClassExpr).getInnerScope()
 }
 
-/** Holds if `f` is a class method. */
-private predicate isClassMethod(Function f) {
-  f.getADecorator() = API::builtin("classmethod").asSource().asExpr()
-  or
-  f.getName() in ["__new__", "__init_subclass__", "__metaclass__", "__class_getitem__"]
-}
-
-/** Holds if `f` is a static method. */
-private predicate isStaticMethod(Function f) {
-  f.getADecorator() = API::builtin("staticmethod").asSource().asExpr()
-}
-
 /** Holds if `c` is a Zope interface. */
 private predicate isZopeInterface(Class c) {
   c =
@@ -55,8 +46,8 @@ private predicate usedInInit(Function f, Class c) {
 /** Holds if the first parameter of `f` should be named `self`. */
 predicate shouldBeSelf(Function f, Class c) {
   methodOfClass(f, c) and
-  not isStaticMethod(f) and
-  not isClassMethod(f) and
+  not isStaticmethod(f) and
+  not isClassmethod(f) and
   not isMetaclass(c) and
   not isZopeInterface(c) and
   not usedInInit(f, c)
@@ -65,24 +56,29 @@ predicate shouldBeSelf(Function f, Class c) {
 /** Holds if the first parameter of `f` should be named `cls`. */
 predicate shouldBeCls(Function f, Class c) {
   methodOfClass(f, c) and
-  not isStaticMethod(f) and
+  not isStaticmethod(f) and
   (
-    isClassMethod(f) and not isMetaclass(c)
+    isClassmethod(f) and not isMetaclass(c)
     or
-    isMetaclass(c) and not isClassMethod(f)
+    isMetaclass(c) and not isClassmethod(f)
   )
 }
 
 /** Holds if the first parameter of `f` is named `self`. */
 predicate firstArgNamedSelf(Function f) { f.getArgName(0) = "self" }
 
-/** Holds if the first parameter of `f` is named `cls`. */
-predicate firstArgNamedCls(Function f) {
+/** Holds if the first parameter of `f` refers to the class - it is either named `cls`, or it is named `self` and is a method of a metaclass. */
+predicate firstArgRefersToCls(Function f, Class c) {
+  methodOfClass(f, c) and
   exists(string argname | argname = f.getArgName(0) |
     argname = "cls"
     or
     /* Not PEP8, but relatively common */
     argname = "mcls"
+    or
+    /* If c is a metaclass, allow arguments named `self`. */
+    argname = "self" and
+    isMetaclass(c)
   )
 }
 
@@ -92,17 +88,11 @@ predicate firstArgShouldBeNamedSelfAndIsnt(Function f) {
   not firstArgNamedSelf(f)
 }
 
-/** Holds if `f` is a regular method of a metaclass, and its first argument is named `self`. */
-private predicate metaclassNamedSelf(Function f, Class c) {
-  methodOfClass(f, c) and
-  firstArgNamedSelf(f) and
-  isMetaclass(c) and
-  not isClassMethod(f)
-}
-
 /** Holds if the first parameter of `f` should be named `cls`, but isn't. */
-predicate firstArgShouldBeNamedClsAndIsnt(Function f) {
-  shouldBeCls(f, _) and
-  not firstArgNamedCls(f) and
-  not metaclassNamedSelf(f, _)
+predicate firstArgShouldReferToClsAndDoesnt(Function f) {
+  exists(Class c |
+    methodOfClass(f, c) and
+    shouldBeCls(f, c) and
+    not firstArgRefersToCls(f, c)
+  )
 }

python/ql/src/Functions/NonCls.ql

@@ -16,7 +16,7 @@ import MethodArgNames
 
 from Function f, string message
 where
-  firstArgShouldBeNamedClsAndIsnt(f) and
+  firstArgShouldReferToClsAndDoesnt(f) and
   (
     if exists(f.getArgName(0))
     then

python/ql/test/query-tests/Functions/methodArgNames/parameter_names.py

@@ -141,3 +141,10 @@ class SpecialMethodNames(object):
 
     def __class_getitem__(cls):
         pass
+
+from dataclasses import dataclass, field
+
+@dataclass 
+class A:
+    # Lambdas used in initilisation aren't methods.
+    x: int = field(default_factory = lambda: 2)

python/ql/test/query-tests/Functions/methodArgNames/test.ql

@@ -14,7 +14,7 @@ module MethodArgTest implements TestSig {
         firstArgShouldBeNamedSelfAndIsnt(f) and
         tag = "shouldBeSelf"
         or
-        firstArgShouldBeNamedClsAndIsnt(f) and
+        firstArgShouldReferToClsAndDoesnt(f) and
         tag = "shouldBeCls"
       )
     )