Add basic support for @classmethod. We ascribe the type Any to the first argument to classmethods; at some point this will need to be tightened up to a bespoke class type.

jswalden · jswalden · commit db9efde92f58 · 2014-07-17T12:41:09.000-07:00
diff --git a/mypy/checkmember.py b/mypy/checkmember.py
@@ -47,7 +47,7 @@ def analyse_member_access(name: str, typ: Type, node: Context, is_lvalue: bool,
         method = info.get_method(name)
         if method:
             if is_lvalue:
-                msg.fail(messages.CANNOT_ASSIGN_TO_METHOD, node)
+                msg.cant_assign_to_method(node)
             typ = map_instance_to_supertype(typ, method.info)
             return expand_type_by_instance(method_type(method), typ)
         else:
@@ -107,20 +107,26 @@ def analyse_member_var_access(name: str, itype: Instance, info: TypeInfo,
         itype = map_instance_to_supertype(itype, var.info)
         if var.type:
             t = expand_type_by_instance(var.type, itype)
-            if (var.is_initialized_in_class and isinstance(t, FunctionLike)
-                    and not var.is_staticmethod):
-                # Class-level function object becomes a bound method.
-                functype = cast(FunctionLike, t)
-                check_method_type(functype, itype, node, msg)
-                signature = method_type(functype)
-                if var.is_property:
-                    if is_lvalue:
+            if var.is_initialized_in_class and isinstance(t, FunctionLike):
+                if is_lvalue:
+                    if var.is_property:
                         msg.read_only_property(name, info, node)
-                    # A property cannot have an overloaded type => the cast
-                    # is fine.
-                    return cast(Callable, signature).ret_type
-                else:
-                    return signature
+                    else:
+                        msg.cant_assign_to_method(node)
+
+                if not var.is_staticmethod:
+                    # Class-level function objects and classmethods become bound
+                    # methods: the former to the instance, the latter to the
+                    # class.
+                    functype = cast(FunctionLike, t)
+                    check_method_type(functype, itype, node, msg)
+                    signature = method_type(functype)
+                    if var.is_property:
+                        # A property cannot have an overloaded type => the cast
+                        # is fine.
+                        return cast(Callable, signature).ret_type
+                    else:
+                        return signature
             return t
         else:
             if not var.is_ready:
@@ -166,38 +172,51 @@ def analyse_class_attribute_access(itype: Instance, name: str,
                                    context: Context, is_lvalue: bool,
                                    msg: MessageBuilder) -> Type:
     node = itype.type.get(name)
-    if node:
-        if is_lvalue and isinstance(node.node, FuncDef):
-            msg.fail(messages.CANNOT_ASSIGN_TO_METHOD, context)
-        if is_lvalue and isinstance(node.node, TypeInfo):
-            msg.fail(messages.CANNOT_ASSIGN_TO_TYPE, context)
-        t = node.type
-        if t:
-            return add_class_tvars(t, itype.type)
-        elif isinstance(node.node, TypeInfo):
-            # TODO add second argument
-            return type_object_type(cast(TypeInfo, node.node), None)
-        else:
-            return function_type(cast(FuncBase, node.node))
-    else:
+    if not node:
         return None
 
+    is_decorated = isinstance(node.node, Decorator)
+    is_method = is_decorated or isinstance(node.node, FuncDef)
+    if is_lvalue:
+        if is_method:
+            msg.cant_assign_to_method(context)
+        if isinstance(node.node, TypeInfo):
+            msg.fail(messages.CANNOT_ASSIGN_TO_TYPE, context)
+
+    t = node.type
+    if t:
+        is_classmethod = is_decorated and cast(Decorator, node.node).func.is_class
+        return add_class_tvars(t, itype.type, is_classmethod)
+
+    if isinstance(node.node, TypeInfo):
+        # TODO add second argument
+        return type_object_type(cast(TypeInfo, node.node), None)
+
+    return function_type(cast(FuncBase, node.node))
+
 
-def add_class_tvars(t: Type, info: TypeInfo) -> Type:
+def add_class_tvars(t: Type, info: TypeInfo, is_classmethod: bool) -> Type:
     if isinstance(t, Callable):
         vars = [TypeVarDef(n, i + 1, None)
                 for i, n in enumerate(info.type_vars)]
-        return Callable(t.arg_types,
-                        t.arg_kinds,
-                        t.arg_names,
+        arg_types = t.arg_types
+        arg_kinds = t.arg_kinds
+        arg_names = t.arg_names
+        if is_classmethod:
+            arg_types = arg_types[1:]
+            arg_kinds = arg_kinds[1:]
+            arg_names = arg_names[1:]
+        return Callable(arg_types,
+                        arg_kinds,
+                        arg_names,
                         t.ret_type,
                         t.is_type_obj(),
                         t.name,
                         vars + t.variables,
                         t.bound_vars,
                         t.line, None)
     elif isinstance(t, Overloaded):
-        return Overloaded([cast(Callable, add_class_tvars(i, info))
+        return Overloaded([cast(Callable, add_class_tvars(i, info, is_classmethod))
                            for i in t.items()])
     return t
 
diff --git a/mypy/messages.py b/mypy/messages.py
@@ -565,6 +565,9 @@ def base_class_definitions_incompatible(self, name: str, base1: TypeInfo,
                   'with definition in base class "{}"'.format(
                       name, base1.name(), base2.name()), context)
 
+    def cant_assign_to_method(self, context: Context) -> None:
+        self.fail(CANNOT_ASSIGN_TO_METHOD, context)
+
     def read_only_property(self, name: str, type: TypeInfo,
                            context: Context) -> None:
         self.fail('Property "{}" defined in "{}" is read-only'.format(
diff --git a/mypy/nodes.py b/mypy/nodes.py
@@ -293,6 +293,7 @@ class FuncDef(FuncItem):
     is_conditional = False             # Defined conditionally (within block)?
     is_abstract = False
     is_static = False
+    is_class = False
     is_property = False
     original_def = None # type: FuncDef  # Original conditional definition
     
@@ -367,6 +368,7 @@ class Var(SymbolNode):
     # Is this initialized explicitly to a non-None value in class body?
     is_initialized_in_class = False
     is_staticmethod = False
+    is_classmethod = False
     is_property = False
     
     def __init__(self, name: str, type: 'mypy.types.Type' = None) -> None:
diff --git a/mypy/semanal.py b/mypy/semanal.py
@@ -64,7 +64,7 @@
 from mypy.types import (
     NoneTyp, Callable, Overloaded, Instance, Type, TypeVar, AnyType,
     FunctionLike, UnboundType, TypeList, ErrorType, TypeVarDef,
-    replace_self_type, TupleType
+    replace_leading_arg_type, TupleType
 )
 from mypy.nodes import function_type, implicit_module_attrs
 from mypy.typeanal import TypeAnalyser, TypeAnalyserPass3
@@ -176,8 +176,10 @@ def visit_func_def(self, defn: FuncDef) -> None:
                     self.fail('Method must have at least one argument', defn)
                 elif defn.type:
                     sig = cast(FunctionLike, defn.type)
-                    defn.type = replace_implicit_self_type(
-                        sig, self_type(self.type))
+                    # TODO: A classmethod's first argument should be more
+                    #       precisely typed than Any.
+                    leading_type = AnyType() if defn.is_class else self_type(self.type)
+                    defn.type = replace_implicit_first_type(sig, leading_type)
 
         if self.is_func_scope() and (not defn.is_decorated and
                                      not defn.is_overload):
@@ -288,9 +290,10 @@ def analyse_function(self, defn: FuncItem) -> None:
             if init_:
                 init_.lvalues[0].accept(self)
         
-        # The first argument of a method is like 'self', though the name could
-        # be different.
-        if is_method and defn.args:
+        # The first argument of a non-static, non-class method is like 'self'
+        # (though the name could be different), having the enclosing class's
+        # instance type.
+        if is_method and not defn.is_static and not defn.is_class and defn.args:
             defn.args[0].is_self = True
         
         defn.body.accept(self)
@@ -933,6 +936,11 @@ def visit_decorator(self, dec: Decorator) -> None:
                 dec.func.is_static = True
                 dec.var.is_staticmethod = True
                 self.check_decorated_function_is_method('staticmethod', dec)
+            elif refers_to_fullname(d, 'builtins.classmethod'):
+                removed.append(i)
+                dec.func.is_class = True
+                dec.var.is_classmethod = True
+                self.check_decorated_function_is_method('classmethod', dec)
             elif refers_to_fullname(d, 'builtins.property'):
                 removed.append(i)
                 dec.func.is_property = True
@@ -1606,17 +1614,17 @@ def self_type(typ: TypeInfo) -> Instance:
 
 
 @overload
-def replace_implicit_self_type(sig: Callable, new: Type) -> Callable:
+def replace_implicit_first_type(sig: Callable, new: Type) -> Callable:
     # We can detect implicit self type by it having no representation.
     if not sig.arg_types[0].repr:
-        return replace_self_type(sig, new)
+        return replace_leading_arg_type(sig, new)
     else:
         return sig
 
 @overload
-def replace_implicit_self_type(sig: FunctionLike, new: Type) -> FunctionLike:
+def replace_implicit_first_type(sig: FunctionLike, new: Type) -> FunctionLike:
     osig = cast(Overloaded, sig)
-    return Overloaded([replace_implicit_self_type(i, new)
+    return Overloaded([replace_implicit_first_type(i, new)
                        for i in osig.items()])
 
 
diff --git a/mypy/strconv.py b/mypy/strconv.py
@@ -108,6 +108,8 @@ def visit_func_def(self, o):
             a.insert(-1, 'Abstract')
         if o.is_static:
             a.insert(-1, 'Static')
+        if o.is_class:
+            a.insert(-1, 'Class')
         if o.is_property:
             a.insert(-1, 'Property')
         return self.dump(a, o)
diff --git a/mypy/test/data/check-classes.test b/mypy/test/data/check-classes.test
@@ -717,6 +717,52 @@ int.from_bytes(b'', '')
 int.from_bytes('', '') # E: Argument 1 to "from_bytes" of "int" has incompatible type "str"; expected "bytes"
 [builtins fixtures/staticmethod.py]
 
+[case testAssignStaticMethodOnInstance]
+import typing
+class A:
+  @staticmethod
+  def f(x: int) -> None: pass
+A().f = A.f # E: Cannot assign to a method
+[builtins fixtures/staticmethod.py]
+
+
+-- Class methods
+-- -------------
+
+
+[case testSimpleClassMethod]
+import typing
+class A:
+  @classmethod
+  def f(cls, x: int) -> None: pass
+A.f(1)
+A().f(1)
+A.f('') # E: Argument 1 to "f" of "A" has incompatible type "str"; expected "int"
+A().f('') # E: Argument 1 to "f" of "A" has incompatible type "str"; expected "int"
+[builtins fixtures/classmethod.py]
+
+[case testBuiltinClassMethod]
+import typing
+int.from_bytes(b'', '')
+int.from_bytes('', '') # E: Argument 1 to "from_bytes" of "int" has incompatible type "str"; expected "bytes"
+[builtins fixtures/classmethod.py]
+
+[case testAssignClassMethodOnClass]
+import typing
+class A:
+  @classmethod
+  def f(cls, x: int) -> None: pass
+A.f = A.f # E: Cannot assign to a method
+[builtins fixtures/classmethod.py]
+
+[case testAssignClassMethodOnInstance]
+import typing
+class A:
+  @classmethod
+  def f(cls, x: int) -> None: pass
+A().f = A.f # E: Cannot assign to a method
+[builtins fixtures/classmethod.py]
+
 
 -- Properties
 -- ----------
diff --git a/mypy/test/data/fixtures/classmethod.py b/mypy/test/data/fixtures/classmethod.py
@@ -0,0 +1,16 @@
+import typing
+
+class object:
+    def __init__(self) -> None: pass
+
+class type:
+    def __init__(self, x) -> None: pass
+
+classmethod = object() # Dummy definition.
+
+class int:
+    @classmethod
+    def from_bytes(cls, bytes: bytes, byteorder: str) -> int: pass
+
+class str: pass
+class bytes: pass
diff --git a/mypy/test/data/pythoneval.test b/mypy/test/data/pythoneval.test
@@ -421,6 +421,17 @@ print(A().f('34'))
 12
 34
 
+[case testClassmethod]
+import typing
+class A:
+    @classmethod
+    def f(cls, x: str) -> int: return int(x)
+print(A.f('12'))
+print(A().f('34'))
+[out]
+12
+34
+
 [case testIntMethods]
 import typing
 print(int.from_bytes(b'ab', 'big'))
diff --git a/mypy/test/data/semanal-classes.test b/mypy/test/data/semanal-classes.test
@@ -425,6 +425,47 @@ MypyFile:1(
         Block:3(
           PassStmt:3())))))
 
+[case testClassMethod]
+class A:
+  @classmethod
+  def f(cls, z: int) -> str: pass
+[builtins fixtures/classmethod.py]
+[out]
+MypyFile:1(
+  ClassDef:1(
+    A
+    Decorator:2(
+      Var(f)
+      FuncDef:3(
+        f
+        Args(
+          Var(cls)
+          Var(z))
+        def (cls: Any, z: builtins.int) -> builtins.str
+        Class
+        Block:3(
+          PassStmt:3())))))
+
+[case testClassMethodWithNoArgs]
+class A:
+  @classmethod
+  def f(cls) -> str: pass
+[builtins fixtures/classmethod.py]
+[out]
+MypyFile:1(
+  ClassDef:1(
+    A
+    Decorator:2(
+      Var(f)
+      FuncDef:3(
+        f
+        Args(
+          Var(cls))
+        def (cls: Any) -> builtins.str
+        Class
+        Block:3(
+          PassStmt:3())))))
+
 [case testProperty]
 import typing
 class A:
diff --git a/mypy/test/data/semanal-errors.test b/mypy/test/data/semanal-errors.test
@@ -1023,6 +1023,20 @@ main, line 2: 'staticmethod' used with a non-method
 main: In function "g":
 main, line 6: 'staticmethod' used with a non-method
 
+[case testClassmethodAndNonMethod]
+import typing
+@classmethod
+def f(): pass
+class A:
+  def g(self):
+    @classmethod
+    def h(): pass
+[builtins fixtures/classmethod.py]
+[out]
+main, line 2: 'classmethod' used with a non-method
+main: In function "g":
+main, line 6: 'classmethod' used with a non-method
+
 [case testNonMethodProperty]
 import typing
 @property  # E: 'property' used with a non-method
diff --git a/mypy/transutil.py b/mypy/transutil.py
diff --git a/mypy/treetransform.py b/mypy/treetransform.py
diff --git a/mypy/types.py b/mypy/types.py
diff --git a/stubs/2.7/builtins.py b/stubs/2.7/builtins.py
diff --git a/stubs/3.2/builtins.py b/stubs/3.2/builtins.py
