Rudimentary support for Enum

mypy/checkmember.py

@@ -220,6 +220,9 @@ def analyse_class_attribute_access(itype: Instance,
         if isinstance(node.node, TypeInfo):
             msg.fail(messages.CANNOT_ASSIGN_TO_TYPE, context)
 
+    if itype.type.is_enum and not (is_lvalue or is_decorated or is_method):
+        return itype
+
     t = node.type
     if t:
         is_classmethod = is_decorated and cast(Decorator, node.node).func.is_class

mypy/nodes.py

@@ -1398,6 +1398,7 @@ class TypeInfo(SymbolNode):
     names = None  # type: SymbolTable       # Names defined directly in this type
     is_abstract = False                    # Does the class have any abstract attributes?
     abstract_attributes = None  # type: List[str]
+    is_enum = False
 
     # Information related to type annotations.
 

mypy/semanal.py

@@ -622,6 +622,8 @@ def analyze_base_classes(self, defn: ClassDef) -> None:
                 defn.base_types.append(base)
             elif not isinstance(base, UnboundType):
                 self.fail('Invalid base class', base_expr)
+            if isinstance(base, Instance):
+                defn.info.is_enum = base.type.fullname() == 'enum.Enum'
         # Add 'object' as implicit base if there is no other base class.
         if (not defn.base_types and defn.fullname != 'builtins.object'):
             obj = self.object_type()

mypy/test/data/pythoneval-enum.test

@@ -0,0 +1,18 @@
+-- Test cases for type checking mypy programs using full stubs and running
+-- using CPython.
+--
+-- These are mostly regression tests -- no attempt is made to make these
+-- complete.
+--
+-- This test file checks Enum
+
+[case testEnum]
+from enum import Enum
+class Medal(Enum):
+    gold = 1
+    silver = 2
+    bronze = 3
+m = Medal.gold
+m = 1
+[out]
+_program.py, line 7: Incompatible types in assignment (expression has type "int", variable has type "Medal")

mypy/test/testpythoneval.py

@@ -27,7 +27,8 @@
 python_eval_files = ['pythoneval.test',
                     'python2eval.test']
 
-python_34_eval_files = ['pythoneval-asyncio.test']
+python_34_eval_files = ['pythoneval-asyncio.test',
+                        'pythoneval-enum.test']
 
 # Path to Python 3 interpreter
 python3_path = sys.executable

stubs/3.4/enum.pyi

@@ -0,0 +1,24 @@
+from typing import List, Any, TypeVar
+
+
+class Enum:
+    def __new__(cls, value: Any) -> None: ...
+    def __repr__(self) -> str: ...
+    def __str__(self) -> str: ...
+    def __dir__(self) -> List[str]: ...
+    def __format__(self, format_spec: str) -> str: ...
+    def __hash__(self) -> Any: ...
+    def __reduce_ex__(self, proto: Any) -> Any: ...
+
+    def name(self) -> str: ...
+    def value(self) -> Any: ...
+
+
+class IntEnum(int, Enum):
+    pass
+
+
+_T = TypeVar('_T')
+
+def unique(enumeration: _T) -> _T:
+    pass