From 5518098ad878f51b887c200e6f65131a436495fa Mon Sep 17 00:00:00 2001
From: Andrew Hilger <ahilger@meta.com>
Date: Mon, 21 Oct 2024 07:49:12 -0700
Subject: [PATCH] use python Enum in types.pyi

Summary: After aligning the `py3.types.Enum` and `python.types.Enum` type stub, remove `py3.types.Enum` and replace it with `python.types.Enum` type stub via re-export. This reflects the reality that at runtime, py3 enums now have the base class  `python.types.Enum`. And same for `Flag` and `BadEnum`.

Differential Revision: D64160660

fbshipit-source-id: 9c56e1eb2f7c15bdba8aa7ceb8663a30decf663f
---
 thrift/lib/py3/test/auto_migrate/converter.py |  3 +-
 thrift/lib/py3/test/auto_migrate/enums.py     |  2 +-
 thrift/lib/py3/types.pyi                      | 49 +++----------------
 thrift/test/py/TestServer.py                  |  2 -
 4 files changed, 10 insertions(+), 46 deletions(-)

diff --git a/thrift/lib/py3/test/auto_migrate/converter.py b/thrift/lib/py3/test/auto_migrate/converter.py
index 378ef489946..bf220cccc0c 100644
--- a/thrift/lib/py3/test/auto_migrate/converter.py
+++ b/thrift/lib/py3/test/auto_migrate/converter.py
@@ -16,6 +16,7 @@
 # pyre-strict
 
 import unittest
+from typing import cast
 
 import convertible.thrift_types as python_types
 import convertible.ttypes as py_deprecated_types
@@ -199,7 +200,7 @@ def test_py_bad_enum(self) -> None:
         )._to_py3()
         self.assertIsInstance(simple.color, BadEnum)
         self.assertEqual(
-            simple.color.enum,
+            cast(BadEnum, simple.color).enum,
             py3_types.Color,
         )
 
diff --git a/thrift/lib/py3/test/auto_migrate/enums.py b/thrift/lib/py3/test/auto_migrate/enums.py
index 7401c749c7b..cf97f04ebde 100644
--- a/thrift/lib/py3/test/auto_migrate/enums.py
+++ b/thrift/lib/py3/test/auto_migrate/enums.py
@@ -164,7 +164,7 @@ def test_bad_enum_in_map_items(self) -> None:
     def assertBadEnum(self, e: BadEnum | Enum, cls: Type[_E], val: int) -> None:
         self.assertIsInstance(e, BadEnum)
         self.assertEqual(e.value, val)
-        self.assertEqual(e.enum, cls)
+        self.assertEqual(cast(BadEnum, e).enum, cls)
         self.assertEqual(int(e), val)
 
     def test_pickle(self) -> None:
diff --git a/thrift/lib/py3/types.pyi b/thrift/lib/py3/types.pyi
index b9cd6eb74dc..de757d9e9ef 100644
--- a/thrift/lib/py3/types.pyi
+++ b/thrift/lib/py3/types.pyi
@@ -22,15 +22,21 @@ from typing import (
     Optional,
     Protocol,
     Sequence,
-    SupportsInt,
     Tuple,
     Type,
     TypeVar,
     Union as tUnion,
 )
 
+import thrift.python.types
+
 from thrift.py3.exceptions import GeneratedError
 
+Enum = thrift.python.types.Enum
+EnumMeta = thrift.python.types.EnumMeta
+Flag = thrift.python.types.Flag
+BadEnum = thrift.python.types.BadEnum
+
 _T = TypeVar("_T")
 eT = TypeVar("eT", bound=Enum)
 
@@ -70,44 +76,3 @@ class Container:
 class List(Container): ...
 class Set(Container): ...
 class Map(Container): ...
-
-class EnumMeta(type):
-    def __iter__(cls: Type[_T]) -> Iterator[_T]: ...
-    def __reversed__(cls: Type[_T]) -> Iterator[_T]: ...
-    def __contains__(cls: Type[_T], item: object) -> bool: ...
-    def __getitem__(cls: Type[_T], name: str) -> _T: ...
-    def __len__(cls) -> int: ...
-    @property
-    def __members__(self: Type[_T]) -> Mapping[str, _T]: ...
-
-class Enum(metaclass=EnumMeta):
-    name: Final[str]
-    value: Final[int]
-    def __getattr__(self: eT, name: str) -> eT: ...
-    def __init__(self: eT, value: tUnion[eT, int]) -> None: ...  # __call__ for meta
-    def __repr__(self) -> str: ...
-    def __str__(self) -> str: ...
-    def __hash__(self) -> int: ...
-    def __int__(self) -> int: ...
-    def __index__(self) -> int: ...
-    def __eq__(self, other: object) -> bool: ...
-
-class Flag(Enum):
-    # pyre-fixme[14]: `__contains__` overrides method defined in `EnumMeta`
-    #  inconsistently.
-    def __contains__(self: eT, other: eT) -> bool: ...
-    def __bool__(self) -> bool: ...
-    # pyre-ignore[15]: This is a pyre bug ignore
-    def __or__(self: eT, other: eT) -> eT: ...
-    def __and__(self: eT, other: eT) -> eT: ...
-    def __xor__(self: eT, other: eT) -> eT: ...
-    def __invert__(self: eT, other: eT) -> eT: ...
-
-class BadEnum(SupportsInt):
-    name: str
-    value: int
-    enum: Enum
-    def __init__(self, the_enum: Type[eT], value: int) -> None: ...
-    def __repr__(self) -> str: ...
-    def __int__(self) -> int: ...
-    def __index__(self) -> int: ...
diff --git a/thrift/test/py/TestServer.py b/thrift/test/py/TestServer.py
index c82ffb8f05d..b9b187235d9 100755
--- a/thrift/test/py/TestServer.py
+++ b/thrift/test/py/TestServer.py
@@ -295,8 +295,6 @@ def main() -> None:
                 ThriftTest.ContextProcessor(TestContextHandler(options.port)),
             )
             processor.registerProcessor(
-                # pyre-fixme[16]: Module `SecondService` has no attribute
-                #  `ContextProcessor`.
                 "SecondService",
                 # pyre-fixme[16]: Module `SecondService` has no attribute
                 #  `ContextProcessor`.