From 62c6f6bd79646b13650aa088b2ed2cd4bb9ea08f Mon Sep 17 00:00:00 2001
From: Guido van Rossum <guido@dropbox.com>
Date: Fri, 3 Feb 2017 06:40:51 -0800
Subject: [PATCH 01/10] Support functional API for Enum.

Fixes #2306.
---
 mypy/checker.py                     |  39 ++++----
 mypy/checkexpr.py                   |  25 ++++-
 mypy/nodes.py                       |  19 ++++
 mypy/semanal.py                     | 136 +++++++++++++++++++++++++++-
 mypy/strconv.py                     |   3 +
 mypy/treetransform.py               |   6 +-
 mypy/visitor.py                     |   7 ++
 test-data/unit/pythoneval-enum.test |  88 ++++++++++++++++++
 8 files changed, 299 insertions(+), 24 deletions(-)

diff --git a/mypy/checker.py b/mypy/checker.py
index 4fa03cdf0c7d..312c6082136f 100644
--- a/mypy/checker.py
+++ b/mypy/checker.py
@@ -17,12 +17,17 @@
     TupleExpr, ListExpr, ExpressionStmt, ReturnStmt, IfStmt,
     WhileStmt, OperatorAssignmentStmt, WithStmt, AssertStmt,
     RaiseStmt, TryStmt, ForStmt, DelStmt, CallExpr, IntExpr, StrExpr,
-    UnicodeExpr, OpExpr, UnaryExpr, LambdaExpr, TempNode, SymbolTableNode,
-    Context, Decorator, PrintStmt, LITERAL_TYPE, BreakStmt, PassStmt, ContinueStmt,
-    ComparisonExpr, StarExpr, EllipsisExpr, RefExpr, ImportFrom, ImportAll, ImportBase,
-    ARG_POS, CONTRAVARIANT, COVARIANT, ExecStmt, GlobalDecl, Import, NonlocalDecl,
-    MDEF, Node
-)
+    BytesExpr, UnicodeExpr, FloatExpr, OpExpr, UnaryExpr, CastExpr, RevealTypeExpr, SuperExpr,
+    TypeApplication, DictExpr, SliceExpr, LambdaExpr, TempNode, SymbolTableNode,
+    Context, ListComprehension, ConditionalExpr, GeneratorExpr,
+    Decorator, SetExpr, TypeVarExpr, NewTypeExpr, PrintStmt,
+    LITERAL_TYPE, BreakStmt, PassStmt, ContinueStmt, ComparisonExpr, StarExpr,
+    YieldFromExpr, NamedTupleExpr, TypedDictExpr, SetComprehension,
+    DictionaryComprehension, ComplexExpr, EllipsisExpr, TypeAliasExpr,
+    RefExpr, YieldExpr, BackquoteExpr, Import, ImportFrom, ImportAll, ImportBase,
+    AwaitExpr, PromoteExpr, Node, EnumCallExpr,
+    ARG_POS, MDEF,
+    CONTRAVARIANT, COVARIANT)
 from mypy import nodes
 from mypy.types import (
     Type, AnyType, CallableType, FunctionLike, Overloaded, TupleType, TypedDictType,
@@ -45,7 +50,7 @@
 from mypy.semanal import set_callable_name, refers_to_fullname
 from mypy.erasetype import erase_typevars
 from mypy.expandtype import expand_type, expand_type_by_instance
-from mypy.visitor import StatementVisitor
+from mypy.visitor import NodeVisitor
 from mypy.join import join_types
 from mypy.treetransform import TransformVisitor
 from mypy.binder import ConditionalTypeBinder, get_declaration
@@ -70,7 +75,7 @@
     ])
 
 
-class TypeChecker(StatementVisitor[None]):
+class TypeChecker(NodeVisitor[None]):
     """Mypy type checker.
 
     Type check mypy source files that have been semantically analyzed.
@@ -2259,21 +2264,13 @@ def visit_break_stmt(self, s: BreakStmt) -> None:
 
     def visit_continue_stmt(self, s: ContinueStmt) -> None:
         self.binder.handle_continue()
+        return None
 
-    def visit_exec_stmt(self, s: ExecStmt) -> None:
-        pass
-
-    def visit_global_decl(self, s: GlobalDecl) -> None:
-        pass
-
-    def visit_nonlocal_decl(self, s: NonlocalDecl) -> None:
-        pass
+    def visit_typeddict_expr(self, e: TypedDictExpr) -> Type:
+        return self.expr_checker.visit_typeddict_expr(e)
 
-    def visit_var(self, s: Var) -> None:
-        pass
-
-    def visit_pass_stmt(self, s: PassStmt) -> None:
-        pass
+    def visit_enum_call_expr(self, o: 'mypy.nodes.EnumCallExpr') -> Type:
+        return self.expr_checker.visit_enum_call_expr(o)
 
     #
     # Helpers
diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py
index 51a6f92e9ce6..e9d2f159a53a 100644
--- a/mypy/checkexpr.py
+++ b/mypy/checkexpr.py
@@ -20,7 +20,8 @@
     ConditionalExpr, ComparisonExpr, TempNode, SetComprehension,
     DictionaryComprehension, ComplexExpr, EllipsisExpr, StarExpr, AwaitExpr, YieldExpr,
     YieldFromExpr, TypedDictExpr, PromoteExpr, NewTypeExpr, NamedTupleExpr, TypeVarExpr,
-    TypeAliasExpr, BackquoteExpr, ARG_POS, ARG_NAMED, ARG_STAR, ARG_STAR2, MODULE_REF,
+    TypeAliasExpr, BackquoteExpr, EnumCallExpr,
+    ARG_POS, ARG_NAMED, ARG_STAR, ARG_STAR2, MODULE_REF,
     UNBOUND_TVAR, BOUND_TVAR, LITERAL_TYPE
 )
 from mypy import nodes
@@ -349,6 +350,12 @@ def check_call(self, callee: Type, args: List[Expression],
         """
         arg_messages = arg_messages or self.msg
         if isinstance(callee, CallableType):
+            if (isinstance(callable_node, RefExpr)
+                and callable_node.fullname in ('enum.Enum', 'enum.IntEnum',
+                                               'enum.Flag', 'enum.IntFlag')):
+                # An Enum() call that failed SemanticAnalyzer.check_enum_call().
+                return callee.ret_type, callee
+
             if (callee.is_type_obj() and callee.type_object().is_abstract
                     # Exceptions for Type[...] and classmethod first argument
                     and not callee.from_type_type and not callee.is_classmethod_class):
@@ -2199,6 +2206,22 @@ def visit_namedtuple_expr(self, e: NamedTupleExpr) -> Type:
         # TODO: Perhaps return a type object type?
         return AnyType()
 
+    def visit_enum_call_expr(self, e: EnumCallExpr) -> Type:
+        for name, value in zip(e.items, e.values):
+            if value is not None:
+                typ = self.accept(value)
+                if not isinstance(typ, AnyType):
+                    var = e.info.names[name].node
+                    if isinstance(var, Var):
+                        # Inline TypeCheker.set_inferred_type(),
+                        # without the lvalue.  (This doesn't really do
+                        # much, since the value attribute is defined
+                        # to have type Any in the typeshed stub.)
+                        var.type = typ
+                        var.is_inferred = True
+        # TODO: Perhaps return a type object type?
+        return AnyType()
+
     def visit_typeddict_expr(self, e: TypedDictExpr) -> Type:
         # TODO: Perhaps return a type object type?
         return AnyType()
diff --git a/mypy/nodes.py b/mypy/nodes.py
index e5a26d638518..4584245b9904 100644
--- a/mypy/nodes.py
+++ b/mypy/nodes.py
@@ -1830,6 +1830,25 @@ def accept(self, visitor: ExpressionVisitor[T]) -> T:
         return visitor.visit_typeddict_expr(self)
 
 
+class EnumCallExpr(Expression):
+    """Named tuple expression Enum('name', 'val1 val2 ...')."""
+
+    # The class representation of this enumerated type
+    info = None  # type: TypeInfo
+    # The item names (for debugging)
+    items = None  # type: List[str]
+    values = None  # type: List[Optional[Expression]]
+
+    def __init__(self, info: 'TypeInfo', items: List[str],
+                 values: List[Optional[Expression]]) -> None:
+        self.info = info
+        self.items = items
+        self.values = values
+
+    def accept(self, visitor: ExpressionVisitor[T]) -> T:
+        return visitor.visit_enum_call_expr(self)
+
+
 class PromoteExpr(Expression):
     """Ducktype class decorator expression _promote(...)."""
 
diff --git a/mypy/semanal.py b/mypy/semanal.py
index ea5a350cdfd5..0ff9d427c62c 100644
--- a/mypy/semanal.py
+++ b/mypy/semanal.py
@@ -64,7 +64,7 @@
     YieldFromExpr, NamedTupleExpr, TypedDictExpr, NonlocalDecl, SymbolNode,
     SetComprehension, DictionaryComprehension, TYPE_ALIAS, TypeAliasExpr,
     YieldExpr, ExecStmt, Argument, BackquoteExpr, ImportBase, AwaitExpr,
-    IntExpr, FloatExpr, UnicodeExpr, EllipsisExpr, TempNode,
+    IntExpr, FloatExpr, UnicodeExpr, EllipsisExpr, TempNode, EnumCallExpr,
     COVARIANT, CONTRAVARIANT, INVARIANT, UNBOUND_IMPORTED, LITERAL_YES, ARG_OPT, nongen_builtins,
     collections_type_aliases, get_member_expr_fullname,
 )
@@ -1498,6 +1498,7 @@ def visit_assignment_stmt(self, s: AssignmentStmt) -> None:
         self.process_typevar_declaration(s)
         self.process_namedtuple_definition(s)
         self.process_typeddict_definition(s)
+        self.process_enum_call(s)
 
         if (len(s.lvalues) == 1 and isinstance(s.lvalues[0], NameExpr) and
                 s.lvalues[0].name == '__all__' and s.lvalues[0].kind == GDEF and
@@ -2327,6 +2328,139 @@ def is_classvar(self, typ: Type) -> bool:
     def fail_invalid_classvar(self, context: Context) -> None:
         self.fail('ClassVar can only be used for assignments in class body', context)
 
+    def process_enum_call(self, s: AssignmentStmt) -> None:
+        """Check if s defines an Enum; if yes, store the definition in symbol table."""
+        if len(s.lvalues) != 1 or not isinstance(s.lvalues[0], NameExpr):
+            return
+        lvalue = s.lvalues[0]
+        name = lvalue.name
+        enum_call = self.check_enum_call(s.rvalue, name)
+        if enum_call is None:
+            return
+        # Yes, it's a valid Enum definition. Add it to the symbol table.
+        node = self.lookup(name, s)
+        if node:
+            node.kind = GDEF   # TODO locally defined Enum
+            node.node = enum_call
+
+    def check_enum_call(self, node: Expression, var_name: str = None) -> Optional[TypeInfo]:
+        """Check if a call defines an Enum.
+
+        Example:
+
+          A = enum.Enum('A', 'foo bar')
+
+        is equivalent to:
+
+          class A(enum.Enum):
+              foo = 1
+              bar = 2
+        """
+        if not isinstance(node, CallExpr):
+            return None
+        call = node
+        callee = call.callee
+        if not isinstance(callee, RefExpr):
+            return None
+        fullname = callee.fullname
+        if fullname not in ('enum.Enum', 'enum.IntEnum', 'enum.Flag', 'enum.IntFlag'):
+            return None
+        items, values, ok = self.parse_enum_call_args(call, fullname.split('.')[-1])
+        if not ok:
+            # Error. Construct dummy return value.
+            return self.build_enum_call_typeinfo('Enum', [], fullname)
+        name = cast(StrExpr, call.args[0]).value
+        if name != var_name or self.is_func_scope():
+            # Give it a unique name derived from the line number.
+            name += '@' + str(call.line)
+        info = self.build_enum_call_typeinfo(name, items, fullname)
+        # Store it as a global just in case it would remain anonymous.
+        # (Or in the nearest class if there is one.)
+        stnode = SymbolTableNode(GDEF, info, self.cur_mod_id)
+        if self.type:
+            self.type.names[name] = stnode
+        else:
+            self.globals[name] = stnode
+        call.analyzed = EnumCallExpr(info, items, values)
+        call.analyzed.set_line(call.line, call.column)
+        return info
+
+    def build_enum_call_typeinfo(self, name: str, items: List[str], fullname: str) -> TypeInfo:
+        base = self.named_type_or_none(fullname)
+        assert base is not None
+        info = self.basic_new_typeinfo(name, base)
+        info.is_enum = True
+        for item in items:
+            var = Var(item)
+            var.info = info
+            var.is_property = True
+            info.names[item] = SymbolTableNode(MDEF, var)
+        return info
+
+    def parse_enum_call_args(self, call: CallExpr,
+                             class_name: str) -> Tuple[List[str],
+                                                       List[Optional[Expression]], bool]:
+        args = call.args
+        if len(args) < 2:
+            return self.fail_enum_call_arg("Too few arguments for %s()" % class_name, call)
+        if len(args) > 2:
+            return self.fail_enum_call_arg("Too many arguments for %s()" % class_name, call)
+        if call.arg_kinds != [ARG_POS, ARG_POS]:
+            return self.fail_enum_call_arg("Unexpected arguments to %s()" % class_name, call)
+        if not isinstance(args[0], (StrExpr, UnicodeExpr)):
+            return self.fail_enum_call_arg(
+                "%s() expects a string literal as the first argument" % class_name, call)
+        items = []
+        values = []  # type: List[Optional[Expression]]
+        if isinstance(args[1], (StrExpr, UnicodeExpr)):
+            fields = args[1].value
+            for field in fields.replace(',', ' ').split():
+                items.append(field)
+        elif isinstance(args[1], (TupleExpr, ListExpr)):
+            seq_items = args[1].items
+            if all(isinstance(seq_item, (StrExpr, UnicodeExpr)) for seq_item in seq_items):
+                items = [cast(StrExpr, seq_item).value for seq_item in seq_items]
+            elif all(isinstance(seq_item, (TupleExpr, ListExpr))
+                     and len(seq_item.items) == 2
+                     and isinstance(seq_item.items[0], (StrExpr, UnicodeExpr))
+                     for seq_item in seq_items):
+                for seq_item in seq_items:
+                    assert isinstance(seq_item, (TupleExpr, ListExpr))
+                    name, value = seq_item.items
+                    assert isinstance(name, (StrExpr, UnicodeExpr))
+                    items.append(name.value)
+                    values.append(value)
+            else:
+                return self.fail_enum_call_arg(
+                    "%s() with tuple or list expects strings or (name, value) pairs" %
+                    class_name,
+                    call)
+        elif isinstance(args[1], DictExpr):
+            for key, value in args[1].items:
+                if not isinstance(key, (StrExpr, UnicodeExpr)):
+                    return self.fail_enum_call_arg(
+                        "%s() with dict literal requires string literals" % class_name, call)
+                items.append(key.value)
+                values.append(value)
+        else:
+            # TODO: Allow dict(x=1, y=2) as a substitute for {'x': 1, 'y': 2}?
+            return self.fail_enum_call_arg(
+                "%s() expects a string, tuple, list or dict literal as the second argument" %
+                class_name,
+                call)
+        if len(items) == 0:
+            return self.fail_enum_call_arg("%s() needs at least one item" % class_name, call)
+        if not values:
+            values = [None] * len(items)
+        assert len(items) == len(values)
+        return items, values, True
+
+    def fail_enum_call_arg(self, message: str,
+                           context: Context) -> Tuple[List[str],
+                                                      List[Optional[Expression]], bool]:
+        self.fail(message, context)
+        return [], [], False
+
     def visit_decorator(self, dec: Decorator) -> None:
         for d in dec.decorators:
             d.accept(self)
diff --git a/mypy/strconv.py b/mypy/strconv.py
index 61411f9b13d8..b8bda6d0224c 100644
--- a/mypy/strconv.py
+++ b/mypy/strconv.py
@@ -431,6 +431,9 @@ def visit_namedtuple_expr(self, o: 'mypy.nodes.NamedTupleExpr') -> str:
                                                   o.info.name(),
                                                   o.info.tuple_type)
 
+    def visit_enum_call_expr(self, o: 'mypy.nodes.EnumCallExpr') -> str:
+        return 'EnumCallExpr:{}({}, {})'.format(o.line, o.info.name(), o.items)
+
     def visit_typeddict_expr(self, o: 'mypy.nodes.TypedDictExpr') -> str:
         return 'TypedDictExpr:{}({})'.format(o.line,
                                              o.info.name())
diff --git a/mypy/treetransform.py b/mypy/treetransform.py
index dde17f2ab2f6..f0a3bcba8caa 100644
--- a/mypy/treetransform.py
+++ b/mypy/treetransform.py
@@ -19,7 +19,8 @@
     ComparisonExpr, TempNode, StarExpr, Statement, Expression,
     YieldFromExpr, NamedTupleExpr, TypedDictExpr, NonlocalDecl, SetComprehension,
     DictionaryComprehension, ComplexExpr, TypeAliasExpr, EllipsisExpr,
-    YieldExpr, ExecStmt, Argument, BackquoteExpr, AwaitExpr, OverloadPart
+    YieldExpr, ExecStmt, Argument, BackquoteExpr, AwaitExpr,
+    OverloadPart, EnumCallExpr,
 )
 from mypy.types import Type, FunctionLike
 from mypy.traverser import TraverserVisitor
@@ -486,6 +487,9 @@ def visit_newtype_expr(self, node: NewTypeExpr) -> NewTypeExpr:
     def visit_namedtuple_expr(self, node: NamedTupleExpr) -> NamedTupleExpr:
         return NamedTupleExpr(node.info)
 
+    def visit_enum_call_expr(self, node: EnumCallExpr) -> EnumCallExpr:
+        return EnumCallExpr(node.info, node.items, node.values)
+
     def visit_typeddict_expr(self, node: TypedDictExpr) -> Node:
         return TypedDictExpr(node.info)
 
diff --git a/mypy/visitor.py b/mypy/visitor.py
index ab691b04be34..6bd7520f4fb6 100644
--- a/mypy/visitor.py
+++ b/mypy/visitor.py
@@ -156,6 +156,10 @@ def visit_type_alias_expr(self, o: 'mypy.nodes.TypeAliasExpr') -> T:
     def visit_namedtuple_expr(self, o: 'mypy.nodes.NamedTupleExpr') -> T:
         pass
 
+    @abstractmethod
+    def visit_enum_call_expr(self, o: 'mypy.nodes.EnumCallExpr') -> T:
+        pass
+
     @abstractmethod
     def visit_typeddict_expr(self, o: 'mypy.nodes.TypedDictExpr') -> T:
         pass
@@ -514,6 +518,9 @@ def visit_type_alias_expr(self, o: 'mypy.nodes.TypeAliasExpr') -> T:
     def visit_namedtuple_expr(self, o: 'mypy.nodes.NamedTupleExpr') -> T:
         pass
 
+    def visit_enum_call_expr(self, o: 'mypy.nodes.EnumCallExpr') -> T:
+        pass
+
     def visit_typeddict_expr(self, o: 'mypy.nodes.TypedDictExpr') -> T:
         pass
 
diff --git a/test-data/unit/pythoneval-enum.test b/test-data/unit/pythoneval-enum.test
index dfbd8f73cef0..1f4a34f91f1b 100644
--- a/test-data/unit/pythoneval-enum.test
+++ b/test-data/unit/pythoneval-enum.test
@@ -157,3 +157,91 @@ class E(IntEnum):
 E[1]
 [out]
 _program.py:4: error: Enum index should be a string (actual index type "int")
+
+[case testFunctionalEnumString]
+from enum import Enum, IntEnum
+E = Enum('E', 'foo bar')
+I = IntEnum('I', ' bar, baz ')
+reveal_type(E.foo)
+reveal_type(E.bar.value)
+reveal_type(I.bar)
+reveal_type(I.baz.value)
+[out]
+_program.py:4: error: Revealed type is '_testFunctionalEnumString.E'
+_program.py:5: error: Revealed type is 'Any'
+_program.py:6: error: Revealed type is '_testFunctionalEnumString.I'
+_program.py:7: error: Revealed type is 'builtins.int'
+
+[case testFunctionalEnumListOfStrings]
+from enum import Enum, IntEnum
+E = Enum('E', ('foo', 'bar'))
+F = IntEnum('F', ['bar', 'baz'])
+reveal_type(E.foo)
+reveal_type(F.baz)
+[out]
+_program.py:4: error: Revealed type is '_testFunctionalEnumListOfStrings.E'
+_program.py:5: error: Revealed type is '_testFunctionalEnumListOfStrings.F'
+
+[case testFunctionalEnumListOfPairs]
+from enum import Enum, IntEnum
+E = Enum('E', [('foo', 1), ['bar', 2]])
+F = IntEnum('F', (['bar', 1], ('baz', 2)))
+reveal_type(E.foo)
+reveal_type(F.baz)
+reveal_type(E.foo.value)
+reveal_type(F.bar.name)
+[out]
+_program.py:4: error: Revealed type is '_testFunctionalEnumListOfPairs.E'
+_program.py:5: error: Revealed type is '_testFunctionalEnumListOfPairs.F'
+_program.py:6: error: Revealed type is 'Any'
+_program.py:7: error: Revealed type is 'builtins.str'
+
+[case testFunctionalEnumDict]
+from enum import Enum, IntEnum
+E = Enum('E', {'foo': 1, 'bar': 2})
+F = IntEnum('F', {'bar': 1, 'baz': 2})
+reveal_type(E.foo)
+reveal_type(F.baz)
+reveal_type(E.foo.value)
+reveal_type(F.bar.name)
+[out]
+_program.py:4: error: Revealed type is '_testFunctionalEnumDict.E'
+_program.py:5: error: Revealed type is '_testFunctionalEnumDict.F'
+_program.py:6: error: Revealed type is 'Any'
+_program.py:7: error: Revealed type is 'builtins.str'
+
+[case testFunctionalEnumErrors]
+from enum import Enum, IntEnum
+A = Enum('A')
+B = Enum('B', 42)
+C = Enum('C', 'a b', 'x')
+D = Enum('D', foo)
+bar = 'x y z'
+E = Enum('E', bar)
+I = IntEnum('I')
+J = IntEnum('I', 42)
+K = IntEnum('I', 'p q', 'z')
+L = Enum('L', ' ')
+M = Enum('M', ())
+N = IntEnum('M', [])
+P = Enum('P', [42])
+Q = Enum('Q', [('a', 42, 0)])
+R = IntEnum('R', [[0, 42]])
+S = Enum('S', {1: 1})
+[out]
+_program.py:2: error: Too few arguments for Enum()
+_program.py:3: error: Enum() expects a string, tuple, list or dict literal as the second argument
+_program.py:4: error: Too many arguments for Enum()
+_program.py:5: error: Enum() expects a string, tuple, list or dict literal as the second argument
+_program.py:5: error: Name 'foo' is not defined
+_program.py:7: error: Enum() expects a string, tuple, list or dict literal as the second argument
+_program.py:8: error: Too few arguments for IntEnum()
+_program.py:9: error: IntEnum() expects a string, tuple, list or dict literal as the second argument
+_program.py:10: error: Too many arguments for IntEnum()
+_program.py:11: error: Enum() needs at least one item
+_program.py:12: error: Enum() needs at least one item
+_program.py:13: error: IntEnum() needs at least one item
+_program.py:14: error: Enum() with tuple or list expects strings or (name, value) pairs
+_program.py:15: error: Enum() with tuple or list expects strings or (name, value) pairs
+_program.py:16: error: IntEnum() with tuple or list expects strings or (name, value) pairs
+_program.py:17: error: Enum() with dict literal requires string literals

From e2a9894ed8869264241c1547f750837ae7c332f6 Mon Sep 17 00:00:00 2001
From: Guido van Rossum <guido@dropbox.com>
Date: Fri, 24 Mar 2017 17:17:22 -0700
Subject: [PATCH 02/10] Fix tests

---
 mypy/checker.py | 6 ------
 1 file changed, 6 deletions(-)

diff --git a/mypy/checker.py b/mypy/checker.py
index 312c6082136f..ba99a4b36077 100644
--- a/mypy/checker.py
+++ b/mypy/checker.py
@@ -2266,12 +2266,6 @@ def visit_continue_stmt(self, s: ContinueStmt) -> None:
         self.binder.handle_continue()
         return None
 
-    def visit_typeddict_expr(self, e: TypedDictExpr) -> Type:
-        return self.expr_checker.visit_typeddict_expr(e)
-
-    def visit_enum_call_expr(self, o: 'mypy.nodes.EnumCallExpr') -> Type:
-        return self.expr_checker.visit_enum_call_expr(o)
-
     #
     # Helpers
     #

From 5fe94f15720b4f2d33486aca78dd424a42668fd3 Mon Sep 17 00:00:00 2001
From: Guido van Rossum <guido@dropbox.com>
Date: Mon, 27 Mar 2017 15:26:02 -0700
Subject: [PATCH 03/10] Move Enum tests from runtests to pytest

---
 mypy/test/testcheck.py                        |  1 +
 mypy/test/testpythoneval.py                   |  3 +-
 .../{pythoneval-enum.test => check-enum.test} | 94 +++++++++----------
 test-data/unit/lib-stub/enum.pyi              | 21 +++++
 4 files changed, 68 insertions(+), 51 deletions(-)
 rename test-data/unit/{pythoneval-enum.test => check-enum.test} (54%)
 create mode 100644 test-data/unit/lib-stub/enum.pyi

diff --git a/mypy/test/testcheck.py b/mypy/test/testcheck.py
index 645c44d93fef..77c820e51fa6 100644
--- a/mypy/test/testcheck.py
+++ b/mypy/test/testcheck.py
@@ -73,6 +73,7 @@
     'check-newsyntax.test',
     'check-underscores.test',
     'check-classvar.test',
+    'check-enum.test',
 ]
 
 
diff --git a/mypy/test/testpythoneval.py b/mypy/test/testpythoneval.py
index cc598d226f2f..85e9aa3751d7 100644
--- a/mypy/test/testpythoneval.py
+++ b/mypy/test/testpythoneval.py
@@ -32,8 +32,7 @@
 python_eval_files = ['pythoneval.test',
                      'python2eval.test']
 
-python_34_eval_files = ['pythoneval-asyncio.test',
-                        'pythoneval-enum.test']
+python_34_eval_files = ['pythoneval-asyncio.test']
 
 # Path to Python 3 interpreter
 python3_path = sys.executable
diff --git a/test-data/unit/pythoneval-enum.test b/test-data/unit/check-enum.test
similarity index 54%
rename from test-data/unit/pythoneval-enum.test
rename to test-data/unit/check-enum.test
index 1f4a34f91f1b..c9e19bc2d9ce 100644
--- a/test-data/unit/pythoneval-enum.test
+++ b/test-data/unit/check-enum.test
@@ -1,9 +1,3 @@
--- 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 testEnumBasics]
@@ -15,7 +9,7 @@ class Medal(Enum):
 m = Medal.gold
 m = 1
 [out]
-_program.py:7: error: Incompatible types in assignment (expression has type "int", variable has type "Medal")
+main:7: error: Incompatible types in assignment (expression has type "int", variable has type "Medal")
 
 [case testEnumNameAndValue]
 from enum import Enum
@@ -24,11 +18,12 @@ class Truth(Enum):
     false = False
 x = ''
 x = Truth.true.name
-print(Truth.true.name)
-print(Truth.false.value)
+reveal_type(Truth.true.name)
+reveal_type(Truth.false.value)
+[builtins fixtures/bool.pyi]
 [out]
-true
-False
+main:7: error: Revealed type is 'builtins.str'
+main:8: error: Revealed type is 'Any'
 
 [case testEnumUnique]
 import enum
@@ -39,7 +34,7 @@ class E(enum.Enum):
 x = 1
 x = E.x
 [out]
-_program.py:7: error: Incompatible types in assignment (expression has type "E", variable has type "int")
+main:7: error: Incompatible types in assignment (expression has type "E", variable has type "int")
 
 [case testIntEnum_assignToIntVariable]
 from enum import IntEnum
@@ -51,7 +46,7 @@ n = N.x  # Subclass of int, so it's okay
 s = ''
 s = N.y
 [out]
-_program.py:8: error: Incompatible types in assignment (expression has type "N", variable has type "str")
+main:8: error: Incompatible types in assignment (expression has type "N", variable has type "str")
 
 [case testIntEnum_functionTakingIntEnum]
 from enum import IntEnum
@@ -63,7 +58,7 @@ takes_some_int_enum(SomeIntEnum.x)
 takes_some_int_enum(1)  # Error
 takes_some_int_enum(SomeIntEnum(1))  # How to deal with the above
 [out]
-_program.py:7: error: Argument 1 to "takes_some_int_enum" has incompatible type "int"; expected "SomeIntEnum"
+main:7: error: Argument 1 to "takes_some_int_enum" has incompatible type "int"; expected "SomeIntEnum"
 
 [case testIntEnum_functionTakingInt]
 from enum import IntEnum
@@ -100,9 +95,10 @@ class Color(Enum):
 
 Color.red.m('')
 Color.m2('')
+[builtins fixtures/staticmethod.pyi]
 [out]
-_program.py:11: error: Argument 1 to "m" of "Color" has incompatible type "str"; expected "int"
-_program.py:12: error: Argument 1 to "m2" of "Color" has incompatible type "str"; expected "int"
+main:11: error: Argument 1 to "m" of "Color" has incompatible type "str"; expected "int"
+main:12: error: Argument 1 to "m2" of "Color" has incompatible type "str"; expected "int"
 
 [case testIntEnum_ExtendedIntEnum_functionTakingExtendedIntEnum]
 from enum import IntEnum
@@ -139,7 +135,7 @@ class E(IntEnum):
 x = None  # type: int
 reveal_type(E(x))
 [out]
-_program.py:5: error: Revealed type is '_testEnumCall.E'
+main:5: error: Revealed type is '__main__.E'
 
 [case testEnumIndex]
 from enum import IntEnum
@@ -148,7 +144,7 @@ class E(IntEnum):
 s = None  # type: str
 reveal_type(E[s])
 [out]
-_program.py:5: error: Revealed type is '_testEnumIndex.E'
+main:5: error: Revealed type is '__main__.E'
 
 [case testEnumIndexError]
 from enum import IntEnum
@@ -156,7 +152,7 @@ class E(IntEnum):
     a = 1
 E[1]
 [out]
-_program.py:4: error: Enum index should be a string (actual index type "int")
+main:4: error: Enum index should be a string (actual index type "int")
 
 [case testFunctionalEnumString]
 from enum import Enum, IntEnum
@@ -167,10 +163,10 @@ reveal_type(E.bar.value)
 reveal_type(I.bar)
 reveal_type(I.baz.value)
 [out]
-_program.py:4: error: Revealed type is '_testFunctionalEnumString.E'
-_program.py:5: error: Revealed type is 'Any'
-_program.py:6: error: Revealed type is '_testFunctionalEnumString.I'
-_program.py:7: error: Revealed type is 'builtins.int'
+main:4: error: Revealed type is '__main__.E'
+main:5: error: Revealed type is 'Any'
+main:6: error: Revealed type is '__main__.I'
+main:7: error: Revealed type is 'builtins.int'
 
 [case testFunctionalEnumListOfStrings]
 from enum import Enum, IntEnum
@@ -179,8 +175,8 @@ F = IntEnum('F', ['bar', 'baz'])
 reveal_type(E.foo)
 reveal_type(F.baz)
 [out]
-_program.py:4: error: Revealed type is '_testFunctionalEnumListOfStrings.E'
-_program.py:5: error: Revealed type is '_testFunctionalEnumListOfStrings.F'
+main:4: error: Revealed type is '__main__.E'
+main:5: error: Revealed type is '__main__.F'
 
 [case testFunctionalEnumListOfPairs]
 from enum import Enum, IntEnum
@@ -191,10 +187,10 @@ reveal_type(F.baz)
 reveal_type(E.foo.value)
 reveal_type(F.bar.name)
 [out]
-_program.py:4: error: Revealed type is '_testFunctionalEnumListOfPairs.E'
-_program.py:5: error: Revealed type is '_testFunctionalEnumListOfPairs.F'
-_program.py:6: error: Revealed type is 'Any'
-_program.py:7: error: Revealed type is 'builtins.str'
+main:4: error: Revealed type is '__main__.E'
+main:5: error: Revealed type is '__main__.F'
+main:6: error: Revealed type is 'Any'
+main:7: error: Revealed type is 'builtins.str'
 
 [case testFunctionalEnumDict]
 from enum import Enum, IntEnum
@@ -205,10 +201,10 @@ reveal_type(F.baz)
 reveal_type(E.foo.value)
 reveal_type(F.bar.name)
 [out]
-_program.py:4: error: Revealed type is '_testFunctionalEnumDict.E'
-_program.py:5: error: Revealed type is '_testFunctionalEnumDict.F'
-_program.py:6: error: Revealed type is 'Any'
-_program.py:7: error: Revealed type is 'builtins.str'
+main:4: error: Revealed type is '__main__.E'
+main:5: error: Revealed type is '__main__.F'
+main:6: error: Revealed type is 'Any'
+main:7: error: Revealed type is 'builtins.str'
 
 [case testFunctionalEnumErrors]
 from enum import Enum, IntEnum
@@ -229,19 +225,19 @@ Q = Enum('Q', [('a', 42, 0)])
 R = IntEnum('R', [[0, 42]])
 S = Enum('S', {1: 1})
 [out]
-_program.py:2: error: Too few arguments for Enum()
-_program.py:3: error: Enum() expects a string, tuple, list or dict literal as the second argument
-_program.py:4: error: Too many arguments for Enum()
-_program.py:5: error: Enum() expects a string, tuple, list or dict literal as the second argument
-_program.py:5: error: Name 'foo' is not defined
-_program.py:7: error: Enum() expects a string, tuple, list or dict literal as the second argument
-_program.py:8: error: Too few arguments for IntEnum()
-_program.py:9: error: IntEnum() expects a string, tuple, list or dict literal as the second argument
-_program.py:10: error: Too many arguments for IntEnum()
-_program.py:11: error: Enum() needs at least one item
-_program.py:12: error: Enum() needs at least one item
-_program.py:13: error: IntEnum() needs at least one item
-_program.py:14: error: Enum() with tuple or list expects strings or (name, value) pairs
-_program.py:15: error: Enum() with tuple or list expects strings or (name, value) pairs
-_program.py:16: error: IntEnum() with tuple or list expects strings or (name, value) pairs
-_program.py:17: error: Enum() with dict literal requires string literals
+main:2: error: Too few arguments for Enum()
+main:3: error: Enum() expects a string, tuple, list or dict literal as the second argument
+main:4: error: Too many arguments for Enum()
+main:5: error: Enum() expects a string, tuple, list or dict literal as the second argument
+main:5: error: Name 'foo' is not defined
+main:7: error: Enum() expects a string, tuple, list or dict literal as the second argument
+main:8: error: Too few arguments for IntEnum()
+main:9: error: IntEnum() expects a string, tuple, list or dict literal as the second argument
+main:10: error: Too many arguments for IntEnum()
+main:11: error: Enum() needs at least one item
+main:12: error: Enum() needs at least one item
+main:13: error: IntEnum() needs at least one item
+main:14: error: Enum() with tuple or list expects strings or (name, value) pairs
+main:15: error: Enum() with tuple or list expects strings or (name, value) pairs
+main:16: error: IntEnum() with tuple or list expects strings or (name, value) pairs
+main:17: error: Enum() with dict literal requires string literals
diff --git a/test-data/unit/lib-stub/enum.pyi b/test-data/unit/lib-stub/enum.pyi
new file mode 100644
index 000000000000..5f0e699ef5e2
--- /dev/null
+++ b/test-data/unit/lib-stub/enum.pyi
@@ -0,0 +1,21 @@
+from typing import Any, TypeVar, Union
+
+class Enum:
+    def __new__(cls, value: Any) -> None: pass
+    def __repr__(self) -> str: pass
+    def __str__(self) -> str: pass
+    def __format__(self, format_spec: str) -> str: pass
+    def __hash__(self) -> Any: pass
+    def __reduce_ex__(self, proto: Any) -> Any: pass
+
+    name = ''  # type: str
+    value = None  # type: Any
+
+class IntEnum(int, Enum):
+    value = 0  # type: int
+
+_T = TypeVar('_T')
+
+def unique(enumeration: _T) -> _T: pass
+
+# TODO: Flag, IntFlag?

From 7a8e1bb9f093e4f187764fae706e5fa759d52783 Mon Sep 17 00:00:00 2001
From: Guido van Rossum <guido@dropbox.com>
Date: Mon, 27 Mar 2017 15:57:22 -0700
Subject: [PATCH 04/10] Add some tests for Flag and IntFlag

---
 test-data/unit/check-enum.test   | 32 ++++++++++++++++++++++++++++++++
 test-data/unit/lib-stub/enum.pyi |  9 ++++++++-
 2 files changed, 40 insertions(+), 1 deletion(-)

diff --git a/test-data/unit/check-enum.test b/test-data/unit/check-enum.test
index c9e19bc2d9ce..4b6bfbcc2d06 100644
--- a/test-data/unit/check-enum.test
+++ b/test-data/unit/check-enum.test
@@ -154,6 +154,28 @@ E[1]
 [out]
 main:4: error: Enum index should be a string (actual index type "int")
 
+[case testEnumFlag]
+from enum import Flag
+class C(Flag):
+    a = 1
+    b = 2
+x = C.a
+x = 1
+x = x | C.b
+[out]
+main:6: error: Incompatible types in assignment (expression has type "int", variable has type "C")
+
+[case testEnumIntFlag]
+from enum import IntFlag
+class C(IntFlag):
+    a = 1
+    b = 2
+x = C.a
+x = 1
+x = x | C.b
+[out]
+main:6: error: Incompatible types in assignment (expression has type "int", variable has type "C")
+
 [case testFunctionalEnumString]
 from enum import Enum, IntEnum
 E = Enum('E', 'foo bar')
@@ -241,3 +263,13 @@ main:14: error: Enum() with tuple or list expects strings or (name, value) pairs
 main:15: error: Enum() with tuple or list expects strings or (name, value) pairs
 main:16: error: IntEnum() with tuple or list expects strings or (name, value) pairs
 main:17: error: Enum() with dict literal requires string literals
+
+[case testFunctionalEnumFlag]
+from enum import Flag, IntFlag
+A = Flag('A', 'x y')
+B = IntFlag('B', 'a b')
+reveal_type(A.x)
+reveal_type(B.a)
+[out]
+main:4: error: Revealed type is '__main__.A'
+main:5: error: Revealed type is '__main__.B'
diff --git a/test-data/unit/lib-stub/enum.pyi b/test-data/unit/lib-stub/enum.pyi
index 5f0e699ef5e2..facf519bd5df 100644
--- a/test-data/unit/lib-stub/enum.pyi
+++ b/test-data/unit/lib-stub/enum.pyi
@@ -18,4 +18,11 @@ _T = TypeVar('_T')
 
 def unique(enumeration: _T) -> _T: pass
 
-# TODO: Flag, IntFlag?
+# In reality Flag and IntFlag are 3.6 only
+
+class Flag(Enum):
+    def __or__(self: _T, other: Union[int, _T]) -> _T: pass
+
+
+class IntFlag(int, Flag):
+    def __and__(self: _T, other: Union[int, _T]) -> _T: pass

From 7b0d959900189016c909fbee92e2d2daa4dfa632 Mon Sep 17 00:00:00 2001
From: Guido van Rossum <guido@dropbox.com>
Date: Mon, 27 Mar 2017 15:59:55 -0700
Subject: [PATCH 05/10] Tests for anonymous Enum and for Enum in class body.

---
 test-data/unit/check-enum.test | 49 ++++++++++++++++++++++++++++++++++
 1 file changed, 49 insertions(+)

diff --git a/test-data/unit/check-enum.test b/test-data/unit/check-enum.test
index 4b6bfbcc2d06..8d4f24554b13 100644
--- a/test-data/unit/check-enum.test
+++ b/test-data/unit/check-enum.test
@@ -176,6 +176,32 @@ x = x | C.b
 [out]
 main:6: error: Incompatible types in assignment (expression has type "int", variable has type "C")
 
+[case testAnonymousEnum]
+from enum import Enum
+class A:
+    def f(self) -> None:
+        class E(Enum):
+            a = 1
+        self.x = E.a
+a = A()
+reveal_type(a.x)
+[out]
+main:7: error: Revealed type is '__main__.A.E@4'
+
+[case testEnumInClassBody]
+from enum import Enum
+class A:
+    class E(Enum):
+        a = 1
+class B:
+    class E(Enum):
+        a = 1
+x = A.E.a
+y = B.E.a
+x = y
+[out]
+main:10: error: Incompatible types in assignment (expression has type "__main__.B.E", variable has type "__main__.A.E")
+
 [case testFunctionalEnumString]
 from enum import Enum, IntEnum
 E = Enum('E', 'foo bar')
@@ -273,3 +299,26 @@ reveal_type(B.a)
 [out]
 main:4: error: Revealed type is '__main__.A'
 main:5: error: Revealed type is '__main__.B'
+
+[case testAnonymousFunctionalEnum]
+from enum import Enum
+class A:
+    def f(self) -> None:
+        E = Enum('E', 'a b')
+        self.x = E.a
+a = A()
+reveal_type(a.x)
+[out]
+main:7: error: Revealed type is '__main__.A.E@4'
+
+[case testFunctionalEnumInClassBody]
+from enum import Enum
+class A:
+    E = Enum('E', 'a b')
+class B:
+    E = Enum('E', 'a b')
+x = A.E.a
+y = B.E.a
+x = y
+[out]
+main:8: error: Incompatible types in assignment (expression has type "__main__.B.E", variable has type "__main__.A.E")

From 63ae048f88b58e1c8404ea1029411b8fa27edf27 Mon Sep 17 00:00:00 2001
From: Guido van Rossum <guido@dropbox.com>
Date: Mon, 27 Mar 2017 16:31:02 -0700
Subject: [PATCH 06/10] Test errors for unexpected arg kinds

---
 test-data/unit/check-enum.test | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/test-data/unit/check-enum.test b/test-data/unit/check-enum.test
index 8d4f24554b13..9c80e050e38d 100644
--- a/test-data/unit/check-enum.test
+++ b/test-data/unit/check-enum.test
@@ -272,6 +272,9 @@ P = Enum('P', [42])
 Q = Enum('Q', [('a', 42, 0)])
 R = IntEnum('R', [[0, 42]])
 S = Enum('S', {1: 1})
+T = Enum('T', keyword='a b')
+U = Enum('U', *['a'])
+V = Enum('U', **{'a': 1})
 [out]
 main:2: error: Too few arguments for Enum()
 main:3: error: Enum() expects a string, tuple, list or dict literal as the second argument
@@ -289,6 +292,9 @@ main:14: error: Enum() with tuple or list expects strings or (name, value) pairs
 main:15: error: Enum() with tuple or list expects strings or (name, value) pairs
 main:16: error: IntEnum() with tuple or list expects strings or (name, value) pairs
 main:17: error: Enum() with dict literal requires string literals
+main:18: error: Unexpected arguments to Enum()
+main:19: error: Unexpected arguments to Enum()
+main:20: error: Unexpected arguments to Enum()
 
 [case testFunctionalEnumFlag]
 from enum import Flag, IntFlag

From 7fed2ece757d399d80cf0865e3a83d43716ffaff Mon Sep 17 00:00:00 2001
From: Guido van Rossum <guido@dropbox.com>
Date: Mon, 27 Mar 2017 16:32:11 -0700
Subject: [PATCH 07/10] Test missing Enum attribute

---
 test-data/unit/check-enum.test | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/test-data/unit/check-enum.test b/test-data/unit/check-enum.test
index 9c80e050e38d..e3e303758a99 100644
--- a/test-data/unit/check-enum.test
+++ b/test-data/unit/check-enum.test
@@ -275,6 +275,8 @@ S = Enum('S', {1: 1})
 T = Enum('T', keyword='a b')
 U = Enum('U', *['a'])
 V = Enum('U', **{'a': 1})
+W = Enum('W', 'a b')
+W.c
 [out]
 main:2: error: Too few arguments for Enum()
 main:3: error: Enum() expects a string, tuple, list or dict literal as the second argument
@@ -295,6 +297,7 @@ main:17: error: Enum() with dict literal requires string literals
 main:18: error: Unexpected arguments to Enum()
 main:19: error: Unexpected arguments to Enum()
 main:20: error: Unexpected arguments to Enum()
+main:22: error: "W" has no attribute "c"
 
 [case testFunctionalEnumFlag]
 from enum import Flag, IntFlag

From f624fdf2158bc034b73398c15e37e71bf861e4eb Mon Sep 17 00:00:00 2001
From: Guido van Rossum <guido@dropbox.com>
Date: Mon, 27 Mar 2017 16:34:11 -0700
Subject: [PATCH 08/10] Test python2 with unicode, bytes

---
 test-data/unit/check-enum.test | 16 ++++++++++++++++
 1 file changed, 16 insertions(+)

diff --git a/test-data/unit/check-enum.test b/test-data/unit/check-enum.test
index e3e303758a99..7054f050357c 100644
--- a/test-data/unit/check-enum.test
+++ b/test-data/unit/check-enum.test
@@ -331,3 +331,19 @@ y = B.E.a
 x = y
 [out]
 main:8: error: Incompatible types in assignment (expression has type "__main__.B.E", variable has type "__main__.A.E")
+
+[case testFunctionalEnum_python2]
+from enum import Enum
+Eu = Enum(u'Eu', u'a b')
+Eb = Enum(b'Eb', b'a b')
+Gu = Enum(u'Gu', {u'a': 1})
+Gb = Enum(b'Gb', {b'a': 1})
+Hu = Enum(u'Hu', [u'a'])
+Hb = Enum(b'Hb', [b'a'])
+Eu.a
+Eb.a
+Gu.a
+Gb.a
+Hu.a
+Hb.a
+[out]

From ad0ba2a12818426bd77b34fe61dc53847b0f619b Mon Sep 17 00:00:00 2001
From: Guido van Rossum <guido@dropbox.com>
Date: Mon, 27 Mar 2017 16:44:25 -0700
Subject: [PATCH 09/10] Add minimal incremental test

---
 test-data/unit/check-enum.test | 20 ++++++++++++++++++++
 1 file changed, 20 insertions(+)

diff --git a/test-data/unit/check-enum.test b/test-data/unit/check-enum.test
index 7054f050357c..4f3627089e33 100644
--- a/test-data/unit/check-enum.test
+++ b/test-data/unit/check-enum.test
@@ -347,3 +347,23 @@ Gb.a
 Hu.a
 Hb.a
 [out]
+
+[case testEnumIncremental]
+import m
+reveal_type(m.E.a)
+reveal_type(m.F.b)
+[file m.py]
+from enum import Enum
+class E(Enum):
+    a = 1
+    b = 2
+F = Enum('F', 'a b')
+[rechecked]
+[stale]
+[out1]
+main:2: error: Revealed type is 'm.E'
+main:3: error: Revealed type is 'm.F'
+[out2]
+main:2: error: Revealed type is 'm.E'
+main:3: error: Revealed type is 'm.F'
+

From 9476383b0161be49c16eecb9e1b5ad9b22eb6c74 Mon Sep 17 00:00:00 2001
From: Guido van Rossum <guido@dropbox.com>
Date: Mon, 27 Mar 2017 17:33:10 -0700
Subject: [PATCH 10/10] Fix test output for classig anonymous enum in method

---
 test-data/unit/check-enum.test | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test-data/unit/check-enum.test b/test-data/unit/check-enum.test
index 4f3627089e33..630ab9645243 100644
--- a/test-data/unit/check-enum.test
+++ b/test-data/unit/check-enum.test
@@ -186,7 +186,7 @@ class A:
 a = A()
 reveal_type(a.x)
 [out]
-main:7: error: Revealed type is '__main__.A.E@4'
+main:8: error: Revealed type is '__main__.E@4'
 
 [case testEnumInClassBody]
 from enum import Enum