Repository: python/mypy
Pull request: Add support for classmethods and staticmethods in add_method
### Description

Fixes #7301

Adds `is_classmethod` and `is_staticmethod` options to `add_method` so plugins can create class/static methods.
Includes tests.

The old PR #7796 didn't have tests and is stale for a long time.

## Test Plan
`pytest -k testAddClassMethodPlugin`
`pytest -k testIncrementalAddClassMethodPlugin`


Problem statement:
Plugin add_classmethod
**Feature request:**
A function like `mypy.plugins.common.add_method`, but with support for classmethods.

-----

More context:

I am working on a plugin, and would like to be able to detect and raise errors when an inherited (factory) classmethod is called, but *only* if certain conditions hold for the type (which I am computing in the `base_model_hook` and storing on `ctx.cls.info.metadata`).

The method is inherited from a parent class, and when the method hook is called it only contains references to the parent class, not the calling class (on which I've stored the metadata). (The parent does not have the appropriate type-specific metadata I want to access.)

In order to get around this, I had the idea to directly add a classmethod to the subclass so that it could trigger the method_hook itself. (Note: I'm not sure whether this would work since I'd be generating the method in the `base_class_hook`; I'm not sure whether it would get picked up by the `get_method_hook` or not.)

**(Any insight into how I can create a conditionally-erroring classmethod, where the error occurs based on data computed in the `base_model_hook`, would be appreciated.)**

Whether or not the above approach would work for this purpose, I didn't see an easy way to add a classmethod (similarly to how `add_method` works), and I think this might be a nice capability to suppport.

(This came up in a discussion with @ilevkivskyi on gitter, and he suggested I create an issue.)

Hints from issue discussion:
I think we can just add an optional argument `is_classmethod=False` (maybe also `is_staticmethod=False`?) to the existing function (and essentially copy the corresponding part from `semanal_namedtuple.py`).
Would you like to submit a PR, @dmontagu 
@msullivan I can take a look tonight and give it a shot, but am still very new to mypy, at least from the developer side. So as long as you don't mind it being an "attempt", I can open a PR. (But I haven't put any effort into this yet if someone else is able to do it very quickly/easily.)

Commit message: Update test-data/unit/check-incremental.test

Patch to review:
diff --git a/mypy/plugins/common.py b/mypy/plugins/common.py
index 33e2eaa83006..ab19fc395158 100644
--- a/mypy/plugins/common.py
+++ b/mypy/plugins/common.py
@@ -11,6 +11,7 @@
     Block,
     CallExpr,
     ClassDef,
+    Decorator,
     Expression,
     FuncDef,
     JsonDict,
@@ -28,6 +29,7 @@
     CallableType,
     Overloaded,
     Type,
+    TypeType,
     TypeVarType,
     deserialize_type,
     get_proper_type,
@@ -104,6 +106,8 @@ def add_method(
     return_type: Type,
     self_type: Optional[Type] = None,
     tvar_def: Optional[TypeVarType] = None,
+    is_classmethod: bool = False,
+    is_staticmethod: bool = False,
 ) -> None:
     """
     Adds a new method to a class.
@@ -117,6 +121,8 @@ def add_method(
         return_type=return_type,
         self_type=self_type,
         tvar_def=tvar_def,
+        is_classmethod=is_classmethod,
+        is_staticmethod=is_staticmethod,
     )
 
 
@@ -128,8 +134,15 @@ def add_method_to_class(
     return_type: Type,
     self_type: Optional[Type] = None,
     tvar_def: Optional[TypeVarType] = None,
+    is_classmethod: bool = False,
+    is_staticmethod: bool = False,
 ) -> None:
     """Adds a new method to a class definition."""
+
+    assert not (
+        is_classmethod is True and is_staticmethod is True
+    ), "Can't add a new method that's both staticmethod and classmethod."
+
     info = cls.info
 
     # First remove any previously generated methods with the same name
@@ -139,13 +152,21 @@ def add_method_to_class(
         if sym.plugin_generated and isinstance(sym.node, FuncDef):
             cls.defs.body.remove(sym.node)
 
-    self_type = self_type or fill_typevars(info)
     if isinstance(api, SemanticAnalyzerPluginInterface):
         function_type = api.named_type("builtins.function")
     else:
         function_type = api.named_generic_type("builtins.function", [])
 
-    args = [Argument(Var("self"), self_type, None, ARG_POS)] + args
+    if is_classmethod:
+        self_type = self_type or TypeType(fill_typevars(info))
+        first = [Argument(Var("_cls"), self_type, None, ARG_POS, True)]
+    elif is_staticmethod:
+        first = []
+    else:
+        self_type = self_type or fill_typevars(info)
+        first = [Argument(Var("self"), self_type, None, ARG_POS)]
+    args = first + args
+
     arg_types, arg_names, arg_kinds = [], [], []
     for arg in args:
         assert arg.type_annotation, "All arguments must be fully typed."
@@ -160,6 +181,8 @@ def add_method_to_class(
     func = FuncDef(name, args, Block([PassStmt()]))
     func.info = info
     func.type = set_callable_name(signature, func)
+    func.is_class = is_classmethod
+    func.is_static = is_staticmethod
     func._fullname = info.fullname + "." + name
     func.line = info.line
 
@@ -170,7 +193,21 @@ def add_method_to_class(
         r_name = get_unique_redefinition_name(name, info.names)
         info.names[r_name] = info.names[name]
 
-    info.names[name] = SymbolTableNode(MDEF, func, plugin_generated=True)
+    # Add decorator for is_staticmethod. It's unnecessary for is_classmethod.
+    if is_staticmethod:
+        func.is_decorated = True
+        v = Var(name, func.type)
+        v.info = info
+        v._fullname = func._fullname
+        v.is_staticmethod = True
+        dec = Decorator(func, [], v)
+        dec.line = info.line
+        sym = SymbolTableNode(MDEF, dec)
+    else:
+        sym = SymbolTableNode(MDEF, func)
+    sym.plugin_generated = True
+    info.names[name] = sym
+
     info.defn.defs.body.append(func)
 
 
diff --git a/test-data/unit/check-custom-plugin.test b/test-data/unit/check-custom-plugin.test
index ee19113f000f..a716109d345e 100644
--- a/test-data/unit/check-custom-plugin.test
+++ b/test-data/unit/check-custom-plugin.test
@@ -991,3 +991,17 @@ class Cls:
 [file mypy.ini]
 \[mypy]
 plugins=<ROOT>/test-data/unit/plugins/class_attr_hook.py
+
+[case testAddClassMethodPlugin]
+# flags: --config-file tmp/mypy.ini
+class BaseAddMethod: pass
+
+class MyClass(BaseAddMethod):
+    pass
+
+my_class = MyClass()
+reveal_type(MyClass.foo_classmethod)  # N: Revealed type is "def ()"
+reveal_type(MyClass.foo_staticmethod)  # N: Revealed type is "def (builtins.int) -> builtins.str"
+[file mypy.ini]
+\[mypy]
+plugins=<ROOT>/test-data/unit/plugins/add_classmethod.py
diff --git a/test-data/unit/check-incremental.test b/test-data/unit/check-incremental.test
index 44452e2072b3..3920b279e0fc 100644
--- a/test-data/unit/check-incremental.test
+++ b/test-data/unit/check-incremental.test
@@ -5911,6 +5911,35 @@ tmp/c.py:4: note: Revealed type is "TypedDict('a.N', {'r': Union[TypedDict('b.M'
 tmp/c.py:5: error: Incompatible types in assignment (expression has type "Optional[N]", variable has type "int")
 tmp/c.py:7: note: Revealed type is "TypedDict('a.N', {'r': Union[TypedDict('b.M', {'r': Union[..., None], 'x': builtins.int}), None], 'x': builtins.int})"
 
+[case testIncrementalAddClassMethodPlugin]
+# flags: --config-file tmp/mypy.ini
+import b
+
+[file mypy.ini]
+\[mypy]
+plugins=<ROOT>/test-data/unit/plugins/add_classmethod.py
+
+[file a.py]
+class BaseAddMethod: pass
+
+class MyClass(BaseAddMethod):
+    pass
+
+[file b.py]
+import a
+
+[file b.py.2]
+import a
+
+my_class = a.MyClass()
+reveal_type(MyClass.foo_classmethod)
+reveal_type(MyClass.foo_staticmethod)
+
+[rechecked b]
+[out2]
+tmp/b.py:4: note: Revealed type is "def ()"
+tmp/b.py:5: note:  Revealed type is "def (builtins.int) -> builtins.str"
+
 [case testGenericNamedTupleSerialization]
 import b
 [file a.py]
diff --git a/test-data/unit/plugins/add_classmethod.py b/test-data/unit/plugins/add_classmethod.py
new file mode 100644
index 000000000000..5aacc69a8f01
--- /dev/null
+++ b/test-data/unit/plugins/add_classmethod.py
@@ -0,0 +1,28 @@
+from typing import Callable, Optional
+
+from mypy.nodes import ARG_POS, Argument, Var
+from mypy.plugin import ClassDefContext, Plugin
+from mypy.plugins.common import add_method
+from mypy.types import NoneType
+
+
+class ClassMethodPlugin(Plugin):
+    def get_base_class_hook(self, fullname: str) -> Optional[Callable[[ClassDefContext], None]]:
+        if "BaseAddMethod" in fullname:
+            return add_extra_methods_hook
+        return None
+
+
+def add_extra_methods_hook(ctx: ClassDefContext) -> None:
+    add_method(ctx, "foo_classmethod", [], NoneType(), is_classmethod=True)
+    add_method(
+        ctx,
+        "foo_staticmethod",
+        [Argument(Var(""), ctx.api.named_type("builtins.int"), None, ARG_POS)],
+        ctx.api.named_type("builtins.str"),
+        is_staticmethod=True,
+    )
+
+
+def plugin(version):
+    return ClassMethodPlugin


Code review comments:
diff_hunk: @@ -5911,6 +5911,35 @@ tmp/c.py:4: note: Revealed type is "TypedDict('a.N', {'r': Union[TypedDict('b.M'
 tmp/c.py:5: error: Incompatible types in assignment (expression has type "Optional[N]", variable has type "int")
 tmp/c.py:7: note: Revealed type is "TypedDict('a.N', {'r': Union[TypedDict('b.M', {'r': Union[..., None], 'x': builtins.int}), None], 'x': builtins.int})"
 
+[case testIncrementalAddClassMethodPlugin]
+# flags: --config-file tmp/mypy.ini
+import b
+
+[file mypy.ini]
+\[mypy]
+plugins=<ROOT>/test-data/unit/plugins/add_classmethod.py
+
+[file a.py]
+class BaseAddMethod: pass
+
+class MyClass(BaseAddMethod):
+    pass
+
+[file b.py]
+import a
+
+[file b.py.2]
+import a
+
+my_class = a.MyClass()
+reveal_type(MyClass.foo_classmethod)
+reveal_type(MyClass.foo_staticmethod)
+
+[rechecked b]
+[out2]
+tmp/b.py:4: note: Revealed type is "def ()"
+tmp/b.py:5: note:  Revealed type is "def (builtins.int) -> builtins.str"
+
path: test-data/unit/check-incremental.test
text: @user1:
```suggestion
my_class = a.MyClass()
reveal_type(a.MyClass.foo_classmethod)
reveal_type(a.MyClass.foo_staticmethod)
reveal_type(my_class.foo_classmethod)
reveal_type(my_class.foo_staticmethod)

[rechecked b]
[out2]
tmp/b.py:4: note: Revealed type is "def ()"
tmp/b.py:5: note:  Revealed type is "def (builtins.int) -> builtins.str"
tmp/b.py:6: note: Revealed type is "def ()"
tmp/b.py:7: note:  Revealed type is "def (builtins.int) -> builtins.str"
```

@user1:
Sorry, I didn't get it right the first time :)