Repository: python/mypy (Python)

Issue:
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:
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.)

Patch:
diff --git a/mypy/plugins/common.py b/mypy/plugins/common.py
--- a/mypy/plugins/common.py
+++ b/mypy/plugins/common.py
@@ -9,6 +9,7 @@
     Block,
     CallExpr,
     ClassDef,
+    Decorator,
     Expression,
     FuncDef,
     JsonDict,
@@ -26,6 +27,7 @@
     CallableType,
     Overloaded,
     Type,
+    TypeType,
     TypeVarType,
     deserialize_type,
     get_proper_type,
@@ -102,6 +104,8 @@ def add_method(
     return_type: Type,
     self_type: Type | None = None,
     tvar_def: TypeVarType | None = None,
+    is_classmethod: bool = False,
+    is_staticmethod: bool = False,
 ) -> None:
     """
     Adds a new method to a class.
@@ -115,6 +119,8 @@ def add_method(
         return_type=return_type,
         self_type=self_type,
         tvar_def=tvar_def,
+        is_classmethod=is_classmethod,
+        is_staticmethod=is_staticmethod,
     )
 
 
@@ -126,8 +132,15 @@ def add_method_to_class(
     return_type: Type,
     self_type: Type | None = None,
     tvar_def: TypeVarType | None = 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
@@ -137,13 +150,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."
@@ -158,6 +179,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
 
@@ -168,7 +191,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/plugins/add_classmethod.py b/test-data/unit/plugins/add_classmethod.py
new file mode 100644
--- /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


Test patch:
diff --git a/test-data/unit/check-custom-plugin.test b/test-data/unit/check-custom-plugin.test
--- 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
--- a/test-data/unit/check-incremental.test
+++ b/test-data/unit/check-incremental.test
@@ -5911,6 +5911,38 @@ 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(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"
 [case testGenericNamedTupleSerialization]
 import b
 [file a.py]