align python.Enum with py3.Enum

Summary: Move methods items needed from `py3/types.pyi` to `python/types.pyi` for `EnumMeta`, `Enum`, and `Flag`.

Reviewed By: yoney

Differential Revision: D64160655

fbshipit-source-id: 1ab8171f69d3e1f754b3ce9a9b4f2d2346577958

thrift/compiler/generate/templates/py3/types.pyi.mustache

@@ -52,6 +52,8 @@ class {{enum:name}}(thrift.py3.types.{{^enum:flags?}}Enum{{/enum:flags?}}{{#enum
     def _to_python(self) -> "{{> common/thrift_types_module_path}}.{{enum:name}}": ...   # type: ignore
     def _to_py3(self) -> {{enum:name}}: ...
     def _to_py_deprecated(self) -> int: ...
+    def __int__(self) -> int: ...
+    def __index__(self) -> int: ...
 
 
 {{/program:enums}}

thrift/compiler/test/fixtures/basic-annotations/out/py3/gen-py3/module/types.pyi

@@ -22,6 +22,8 @@ class MyEnum(thrift.py3.types.Enum):
     def _to_python(self) -> "module.thrift_types.MyEnum": ...   # type: ignore
     def _to_py3(self) -> MyEnum: ...
     def _to_py_deprecated(self) -> int: ...
+    def __int__(self) -> int: ...
+    def __index__(self) -> int: ...
 
 
 class MyStructNestedAnnotation(thrift.py3.types.Struct, _typing.Hashable):

thrift/compiler/test/fixtures/basic-enum/out/py3/gen-py3/test/fixtures/enumstrict/module/types.pyi

@@ -19,6 +19,8 @@ class EmptyEnum(thrift.py3.types.Enum):
     def _to_python(self) -> "test.fixtures.enumstrict.module.thrift_types.EmptyEnum": ...   # type: ignore
     def _to_py3(self) -> EmptyEnum: ...
     def _to_py_deprecated(self) -> int: ...
+    def __int__(self) -> int: ...
+    def __index__(self) -> int: ...
 
 
 class MyEnum(thrift.py3.types.Enum):
@@ -27,6 +29,8 @@ class MyEnum(thrift.py3.types.Enum):
     def _to_python(self) -> "test.fixtures.enumstrict.module.thrift_types.MyEnum": ...   # type: ignore
     def _to_py3(self) -> MyEnum: ...
     def _to_py_deprecated(self) -> int: ...
+    def __int__(self) -> int: ...
+    def __index__(self) -> int: ...
 
 
 class MyUseIntrinsicDefaultEnum(thrift.py3.types.Enum):
@@ -36,6 +40,8 @@ class MyUseIntrinsicDefaultEnum(thrift.py3.types.Enum):
     def _to_python(self) -> "test.fixtures.enumstrict.module.thrift_types.MyUseIntrinsicDefaultEnum": ...   # type: ignore
     def _to_py3(self) -> MyUseIntrinsicDefaultEnum: ...
     def _to_py_deprecated(self) -> int: ...
+    def __int__(self) -> int: ...
+    def __index__(self) -> int: ...
 
 
 class MyBigEnum(thrift.py3.types.Enum):
@@ -62,6 +68,8 @@ class MyBigEnum(thrift.py3.types.Enum):
     def _to_python(self) -> "test.fixtures.enumstrict.module.thrift_types.MyBigEnum": ...   # type: ignore
     def _to_py3(self) -> MyBigEnum: ...
     def _to_py_deprecated(self) -> int: ...
+    def __int__(self) -> int: ...
+    def __index__(self) -> int: ...
 
 
 class MyStruct(thrift.py3.types.Struct, _typing.Hashable):

thrift/compiler/test/fixtures/basic-stack-arguments/out/py3/gen-py3/module/types.pyi

@@ -21,6 +21,8 @@ class MyEnum(thrift.py3.types.Enum):
     def _to_python(self) -> "module.thrift_types.MyEnum": ...   # type: ignore
     def _to_py3(self) -> MyEnum: ...
     def _to_py_deprecated(self) -> int: ...
+    def __int__(self) -> int: ...
+    def __index__(self) -> int: ...
 
 
 class MyStruct(thrift.py3.types.Struct, _typing.Hashable):

thrift/compiler/test/fixtures/basic/out/py3/gen-py3/test/fixtures/basic/module/types.pyi

@@ -21,6 +21,8 @@ class MyEnum(thrift.py3.types.Enum):
     def _to_python(self) -> "test.fixtures.basic.module.thrift_types.MyEnum": ...   # type: ignore
     def _to_py3(self) -> MyEnum: ...
     def _to_py_deprecated(self) -> int: ...
+    def __int__(self) -> int: ...
+    def __index__(self) -> int: ...
 
 
 class HackEnum(thrift.py3.types.Enum):
@@ -29,6 +31,8 @@ class HackEnum(thrift.py3.types.Enum):
     def _to_python(self) -> "test.fixtures.basic.module.thrift_types.HackEnum": ...   # type: ignore
     def _to_py3(self) -> HackEnum: ...
     def _to_py_deprecated(self) -> int: ...
+    def __int__(self) -> int: ...
+    def __index__(self) -> int: ...
 
 
 class MyStruct(thrift.py3.types.Struct, _typing.Hashable):

thrift/compiler/test/fixtures/constants/out/py3/gen-py3/module/types.pyi

@@ -19,6 +19,8 @@ class EmptyEnum(thrift.py3.types.Enum):
     def _to_python(self) -> "module.thrift_types.EmptyEnum": ...   # type: ignore
     def _to_py3(self) -> EmptyEnum: ...
     def _to_py_deprecated(self) -> int: ...
+    def __int__(self) -> int: ...
+    def __index__(self) -> int: ...
 
 
 class City(thrift.py3.types.Enum):
@@ -29,6 +31,8 @@ class City(thrift.py3.types.Enum):
     def _to_python(self) -> "module.thrift_types.City": ...   # type: ignore
     def _to_py3(self) -> City: ...
     def _to_py_deprecated(self) -> int: ...
+    def __int__(self) -> int: ...
+    def __index__(self) -> int: ...
 
 
 class Company(thrift.py3.types.Enum):
@@ -39,6 +43,8 @@ class Company(thrift.py3.types.Enum):
     def _to_python(self) -> "module.thrift_types.Company": ...   # type: ignore
     def _to_py3(self) -> Company: ...
     def _to_py_deprecated(self) -> int: ...
+    def __int__(self) -> int: ...
+    def __index__(self) -> int: ...
 
 
 class Internship(thrift.py3.types.Struct, _typing.Hashable):

thrift/compiler/test/fixtures/enums/out/py3/gen-py3/test/fixtures/enums/module/types.pyi

@@ -23,6 +23,8 @@ class Metasyntactic(thrift.py3.types.Enum):
     def _to_python(self) -> "test.fixtures.enums.module.thrift_types.Metasyntactic": ...   # type: ignore
     def _to_py3(self) -> Metasyntactic: ...
     def _to_py_deprecated(self) -> int: ...
+    def __int__(self) -> int: ...
+    def __index__(self) -> int: ...
 
 
 class MyEnum1(thrift.py3.types.Enum):
@@ -35,6 +37,8 @@ class MyEnum1(thrift.py3.types.Enum):
     def _to_python(self) -> "test.fixtures.enums.module.thrift_types.MyEnum1": ...   # type: ignore
     def _to_py3(self) -> MyEnum1: ...
     def _to_py_deprecated(self) -> int: ...
+    def __int__(self) -> int: ...
+    def __index__(self) -> int: ...
 
 
 class MyEnum2(thrift.py3.types.Enum):
@@ -44,6 +48,8 @@ class MyEnum2(thrift.py3.types.Enum):
     def _to_python(self) -> "test.fixtures.enums.module.thrift_types.MyEnum2": ...   # type: ignore
     def _to_py3(self) -> MyEnum2: ...
     def _to_py_deprecated(self) -> int: ...
+    def __int__(self) -> int: ...
+    def __index__(self) -> int: ...
 
 
 class MyEnum3(thrift.py3.types.Enum):
@@ -56,6 +62,8 @@ class MyEnum3(thrift.py3.types.Enum):
     def _to_python(self) -> "test.fixtures.enums.module.thrift_types.MyEnum3": ...   # type: ignore
     def _to_py3(self) -> MyEnum3: ...
     def _to_py_deprecated(self) -> int: ...
+    def __int__(self) -> int: ...
+    def __index__(self) -> int: ...
 
 
 class MyEnum4(thrift.py3.types.Enum):
@@ -66,6 +74,8 @@ class MyEnum4(thrift.py3.types.Enum):
     def _to_python(self) -> "test.fixtures.enums.module.thrift_types.MyEnum4": ...   # type: ignore
     def _to_py3(self) -> MyEnum4: ...
     def _to_py_deprecated(self) -> int: ...
+    def __int__(self) -> int: ...
+    def __index__(self) -> int: ...
 
 
 class MyBitmaskEnum1(thrift.py3.types.Enum):
@@ -75,6 +85,8 @@ class MyBitmaskEnum1(thrift.py3.types.Enum):
     def _to_python(self) -> "test.fixtures.enums.module.thrift_types.MyBitmaskEnum1": ...   # type: ignore
     def _to_py3(self) -> MyBitmaskEnum1: ...
     def _to_py_deprecated(self) -> int: ...
+    def __int__(self) -> int: ...
+    def __index__(self) -> int: ...
 
 
 class MyBitmaskEnum2(thrift.py3.types.Enum):
@@ -84,6 +96,8 @@ class MyBitmaskEnum2(thrift.py3.types.Enum):
     def _to_python(self) -> "test.fixtures.enums.module.thrift_types.MyBitmaskEnum2": ...   # type: ignore
     def _to_py3(self) -> MyBitmaskEnum2: ...
     def _to_py_deprecated(self) -> int: ...
+    def __int__(self) -> int: ...
+    def __index__(self) -> int: ...
 
 
 class SomeStruct(thrift.py3.types.Struct, _typing.Hashable):

thrift/compiler/test/fixtures/mcpp2-compare/out/py3/gen-py3/includes/types.pyi

@@ -21,6 +21,8 @@ class AnEnum(thrift.py3.types.Enum):
     def _to_python(self) -> "includes.thrift_types.AnEnum": ...   # type: ignore
     def _to_py3(self) -> AnEnum: ...
     def _to_py_deprecated(self) -> int: ...
+    def __int__(self) -> int: ...
+    def __index__(self) -> int: ...
 
 
 class AStruct(thrift.py3.types.Struct, _typing.Hashable):

thrift/compiler/test/fixtures/mcpp2-compare/out/py3/gen-py3/module/types.pyi

@@ -23,6 +23,8 @@ class MyEnumA(thrift.py3.types.Enum):
     def _to_python(self) -> "module.thrift_types.MyEnumA": ...   # type: ignore
     def _to_py3(self) -> MyEnumA: ...
     def _to_py_deprecated(self) -> int: ...
+    def __int__(self) -> int: ...
+    def __index__(self) -> int: ...
 
 
 class AnnotatedEnum(thrift.py3.types.Enum):
@@ -32,6 +34,8 @@ class AnnotatedEnum(thrift.py3.types.Enum):
     def _to_python(self) -> "module.thrift_types.AnnotatedEnum": ...   # type: ignore
     def _to_py3(self) -> AnnotatedEnum: ...
     def _to_py_deprecated(self) -> int: ...
+    def __int__(self) -> int: ...
+    def __index__(self) -> int: ...
 
 
 class AnnotatedEnum2(thrift.py3.types.Enum):
@@ -41,13 +45,17 @@ class AnnotatedEnum2(thrift.py3.types.Enum):
     def _to_python(self) -> "module.thrift_types.AnnotatedEnum2": ...   # type: ignore
     def _to_py3(self) -> AnnotatedEnum2: ...
     def _to_py_deprecated(self) -> int: ...
+    def __int__(self) -> int: ...
+    def __index__(self) -> int: ...
 
 
 class MyEnumB(thrift.py3.types.Enum):
     AField: MyEnumB = ...
     def _to_python(self) -> "module.thrift_types.MyEnumB": ...   # type: ignore
     def _to_py3(self) -> MyEnumB: ...
     def _to_py_deprecated(self) -> int: ...
+    def __int__(self) -> int: ...
+    def __index__(self) -> int: ...
 
 
 class Empty(thrift.py3.types.Struct, _typing.Hashable):

thrift/compiler/test/fixtures/optionals/out/py3/gen-py3/module/types.pyi

@@ -22,6 +22,8 @@ class Animal(thrift.py3.types.Enum):
     def _to_python(self) -> "module.thrift_types.Animal": ...   # type: ignore
     def _to_py3(self) -> Animal: ...
     def _to_py_deprecated(self) -> int: ...
+    def __int__(self) -> int: ...
+    def __index__(self) -> int: ...
 
 
 class Color(thrift.py3.types.Struct, _typing.Hashable):

thrift/compiler/test/fixtures/py3/out/py3/gen-py3/module/types.pyi

@@ -25,6 +25,8 @@ class AnEnum(thrift.py3.types.Enum, module.thrift_types._fbthrift_compatible_wit
     def _to_python(self) -> "module.thrift_types.AnEnum": ...   # type: ignore
     def _to_py3(self) -> AnEnum: ...
     def _to_py_deprecated(self) -> int: ...
+    def __int__(self) -> int: ...
+    def __index__(self) -> int: ...
 
 
 class AnEnumRenamed(thrift.py3.types.Enum, module.thrift_types._fbthrift_compatible_with_AnEnumRenamed):
@@ -34,6 +36,8 @@ class AnEnumRenamed(thrift.py3.types.Enum, module.thrift_types._fbthrift_compati
     def _to_python(self) -> "module.thrift_types.AnEnumRenamed": ...   # type: ignore
     def _to_py3(self) -> AnEnumRenamed: ...
     def _to_py_deprecated(self) -> int: ...
+    def __int__(self) -> int: ...
+    def __index__(self) -> int: ...
 
 
 class Flags(thrift.py3.types.Flag, module.thrift_types._fbthrift_compatible_with_Flags):
@@ -44,6 +48,8 @@ class Flags(thrift.py3.types.Flag, module.thrift_types._fbthrift_compatible_with
     def _to_python(self) -> "module.thrift_types.Flags": ...   # type: ignore
     def _to_py3(self) -> Flags: ...
     def _to_py_deprecated(self) -> int: ...
+    def __int__(self) -> int: ...
+    def __index__(self) -> int: ...
 
 
 class SimpleException(thrift.py3.exceptions.GeneratedError, _typing.Hashable, module.thrift_types._fbthrift_compatible_with_SimpleException):

thrift/compiler/test/fixtures/py3/out/py3_auto_migrate/gen-py3/module/types.pyi

@@ -24,6 +24,8 @@ class AnEnum(thrift.py3.types.Enum):
     def _to_python(self) -> "module.thrift_types.AnEnum": ...   # type: ignore
     def _to_py3(self) -> AnEnum: ...
     def _to_py_deprecated(self) -> int: ...
+    def __int__(self) -> int: ...
+    def __index__(self) -> int: ...
 
 
 class AnEnumRenamed(thrift.py3.types.Enum):
@@ -33,6 +35,8 @@ class AnEnumRenamed(thrift.py3.types.Enum):
     def _to_python(self) -> "module.thrift_types.AnEnumRenamed": ...   # type: ignore
     def _to_py3(self) -> AnEnumRenamed: ...
     def _to_py_deprecated(self) -> int: ...
+    def __int__(self) -> int: ...
+    def __index__(self) -> int: ...
 
 
 class Flags(thrift.py3.types.Flag):
@@ -43,6 +47,8 @@ class Flags(thrift.py3.types.Flag):
     def _to_python(self) -> "module.thrift_types.Flags": ...   # type: ignore
     def _to_py3(self) -> Flags: ...
     def _to_py_deprecated(self) -> int: ...
+    def __int__(self) -> int: ...
+    def __index__(self) -> int: ...
 
 
 class SimpleException(thrift.py3.exceptions.GeneratedError, _typing.Hashable):

thrift/compiler/test/fixtures/qualified/out/py3/gen-py3/module0/types.pyi

@@ -22,6 +22,8 @@ class Enum(thrift.py3.types.Enum):
     def _to_python(self) -> "module0.thrift_types.Enum": ...   # type: ignore
     def _to_py3(self) -> Enum: ...
     def _to_py_deprecated(self) -> int: ...
+    def __int__(self) -> int: ...
+    def __index__(self) -> int: ...
 
 
 class Struct(thrift.py3.types.Struct, _typing.Hashable):

thrift/compiler/test/fixtures/qualified/out/py3/gen-py3/module1/types.pyi

@@ -22,6 +22,8 @@ class Enum(thrift.py3.types.Enum):
     def _to_python(self) -> "module1.thrift_types.Enum": ...   # type: ignore
     def _to_py3(self) -> Enum: ...
     def _to_py_deprecated(self) -> int: ...
+    def __int__(self) -> int: ...
+    def __index__(self) -> int: ...
 
 
 class Struct(thrift.py3.types.Struct, _typing.Hashable):

thrift/compiler/test/fixtures/refs/out/py3/gen-py3/module/types.pyi

@@ -21,6 +21,8 @@ class MyEnum(thrift.py3.types.Enum):
     def _to_python(self) -> "module.thrift_types.MyEnum": ...   # type: ignore
     def _to_py3(self) -> MyEnum: ...
     def _to_py_deprecated(self) -> int: ...
+    def __int__(self) -> int: ...
+    def __index__(self) -> int: ...
 
 
 class TypedEnum(thrift.py3.types.Enum):
@@ -29,6 +31,8 @@ class TypedEnum(thrift.py3.types.Enum):
     def _to_python(self) -> "module.thrift_types.TypedEnum": ...   # type: ignore
     def _to_py3(self) -> TypedEnum: ...
     def _to_py_deprecated(self) -> int: ...
+    def __int__(self) -> int: ...
+    def __index__(self) -> int: ...
 
 
 _MyUnionValueType = _typing.Union[None, int, str]

thrift/compiler/test/fixtures/templated-deserialize/out/py3/gen-py3/module/types.pyi

@@ -22,6 +22,8 @@ class MyEnumA(thrift.py3.types.Enum):
     def _to_python(self) -> "module.thrift_types.MyEnumA": ...   # type: ignore
     def _to_py3(self) -> MyEnumA: ...
     def _to_py_deprecated(self) -> int: ...
+    def __int__(self) -> int: ...
+    def __index__(self) -> int: ...
 
 
 class SmallStruct(thrift.py3.types.Struct, _typing.Hashable):

thrift/compiler/test/fixtures/types/out/py3/gen-py3/apache/thrift/fixtures/types/module/types.pyi

@@ -25,6 +25,8 @@ class has_bitwise_ops(thrift.py3.types.Enum):
     def _to_python(self) -> "apache.thrift.fixtures.types.module.thrift_types.has_bitwise_ops": ...   # type: ignore
     def _to_py3(self) -> has_bitwise_ops: ...
     def _to_py_deprecated(self) -> int: ...
+    def __int__(self) -> int: ...
+    def __index__(self) -> int: ...
 
 
 class is_unscoped(thrift.py3.types.Enum):
@@ -33,6 +35,8 @@ class is_unscoped(thrift.py3.types.Enum):
     def _to_python(self) -> "apache.thrift.fixtures.types.module.thrift_types.is_unscoped": ...   # type: ignore
     def _to_py3(self) -> is_unscoped: ...
     def _to_py_deprecated(self) -> int: ...
+    def __int__(self) -> int: ...
+    def __index__(self) -> int: ...
 
 
 class MyForwardRefEnum(thrift.py3.types.Enum):
@@ -41,6 +45,8 @@ class MyForwardRefEnum(thrift.py3.types.Enum):
     def _to_python(self) -> "apache.thrift.fixtures.types.module.thrift_types.MyForwardRefEnum": ...   # type: ignore
     def _to_py3(self) -> MyForwardRefEnum: ...
     def _to_py_deprecated(self) -> int: ...
+    def __int__(self) -> int: ...
+    def __index__(self) -> int: ...
 
 
 class empty_struct(thrift.py3.types.Struct, _typing.Hashable):

thrift/lib/py3/test/auto_migrate/enums.py

@@ -244,6 +244,11 @@ def test_isinstance_Enum(self) -> None:
         self.assertIsInstance(Color.red, Enum)
         self.assertTrue(issubclass(Color, Enum))
 
+    def test_callable(self) -> None:
+        vals = range(3)
+        colors = list(map(Color, vals))
+        self.assertEqual(colors, list(Color))
+
 
 class EnumMetaTests(unittest.TestCase):
     def test_iter_forward_compatible(self) -> None:

thrift/lib/python/test/enums.py

@@ -334,6 +334,12 @@ def test_insinstance_Enum(self) -> None:
         self.assertIsInstance(self.Color.red, Enum)
         self.assertTrue(issubclass(self.Color, Enum))
 
+    def test_callable(self) -> None:
+        vals = range(3)
+        # this is done to verify pyre typestub more than anything
+        colors = list(map(self.Color, vals))
+        self.assertEqual(colors, list(self.Color))
+
 
 # tt = test_types, ser = serializer
 @parameterized_class(