From 708396cf66cbabf435a85a3ac8c40038391a4198 Mon Sep 17 00:00:00 2001
From: Elazar Gershuni <elazarg@gmail.com>
Date: Fri, 20 Apr 2018 00:30:56 +0300
Subject: [PATCH] fix functional enum metaclass

---
 mypy/semanal_enum.py             |  1 +
 test-data/unit/check-enum.test   | 10 ++++++++++
 test-data/unit/lib-stub/enum.pyi | 12 +++++++-----
 3 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/mypy/semanal_enum.py b/mypy/semanal_enum.py
index 664443024ab5..fbfe723b0113 100644
--- a/mypy/semanal_enum.py
+++ b/mypy/semanal_enum.py
@@ -80,6 +80,7 @@ def build_enum_call_typeinfo(self, name: str, items: List[str], fullname: str) -
         base = self.api.named_type_or_none(fullname)
         assert base is not None
         info = self.api.basic_new_typeinfo(name, base)
+        info.metaclass_type = info.calculate_metaclass_type()
         info.is_enum = True
         for item in items:
             var = Var(item)
diff --git a/test-data/unit/check-enum.test b/test-data/unit/check-enum.test
index 450fe8f60a35..c04b1210faa4 100644
--- a/test-data/unit/check-enum.test
+++ b/test-data/unit/check-enum.test
@@ -392,6 +392,16 @@ x = y
 [out]
 main:8: error: Incompatible types in assignment (expression has type "__main__.B.E", variable has type "__main__.A.E")
 
+[case testFunctionalEnumProtocols]
+from enum import IntEnum
+Color = IntEnum('Color', 'red green blue')
+reveal_type(Color['green'])  # E: Revealed type is '__main__.Color'
+for c in Color:
+    reveal_type(c)  # E: Revealed type is '__main__.Color*'
+reveal_type(list(Color))  # E: Revealed type is 'builtins.list[__main__.Color*]'
+
+[builtins fixtures/list.pyi]
+
 [case testEnumWorkWithForward]
 from enum import Enum
 a: E = E.x
diff --git a/test-data/unit/lib-stub/enum.pyi b/test-data/unit/lib-stub/enum.pyi
index 2bef9fc01c17..c49a584bd5a1 100644
--- a/test-data/unit/lib-stub/enum.pyi
+++ b/test-data/unit/lib-stub/enum.pyi
@@ -1,7 +1,11 @@
-from typing import Any, TypeVar, Union
+from typing import Any, TypeVar, Union, Type, Sized, Iterator, Mapping
 
-class EnumMeta(type):
-    pass
+_T = TypeVar('_T')
+
+class EnumMeta(type, Sized):
+    def __iter__(self: Type[_T]) -> Iterator[_T]: pass
+    def __reversed__(self: Type[_T]) -> Iterator[_T]: pass
+    def __getitem__(self: Type[_T], name: str) -> _T: pass
 
 class Enum(metaclass=EnumMeta):
     def __new__(cls, value: Any) -> None: pass
@@ -17,8 +21,6 @@ class Enum(metaclass=EnumMeta):
 class IntEnum(int, Enum):
     value = 0  # type: int
 
-_T = TypeVar('_T')
-
 def unique(enumeration: _T) -> _T: pass
 
 # In reality Flag and IntFlag are 3.6 only