Use unified thrift enums with mutable types

Summary: # TL;DR Provide a single identity for enums in a thrift schema through a .thrift_enums module. # Problem Statement With the introduction of thrift mutable types (`.thrift_mutable_types`) and thrift abstract types (`.thrift_abstract_types`), along with current thrift types (`.thrift_types`), which are immutable there are three different flavors of types for a given thrift schema. Enums are just enums and are neither mutable nor abstract. This implies enums should have a single identity irrespective of whether they are from `.thrift_abstract_types`, `.thrift_types`, or `.thrift_mutable_types`. # Solution To capture this nature of enums, create a single identity for thrift enums via a `.thrift_enums` module and refer to that from the other python flavors. # Current diff: This diff uses the enums from `.thrift_enums` with mutable types in `.thrift_mutable_types`. Reviewed By: ahilger Differential Revision: D66209194 fbshipit-source-id: 464369e222066e16a23d8a2625d9c90f305a1ecb
facebook · Nov 20, 2024 · 8b78f41 · 8b78f41
1 parent 517e16e
commit 8b78f41
Show file tree

Hide file tree

Showing 71 changed files with 199 additions and 1,524 deletions.
diff --git a/thrift/compiler/generate/templates/python/types/types.mustache b/thrift/compiler/generate/templates/python/types/types.mustache
@@ -187,8 +187,24 @@ Until we resolve the above error, register the class as virtual subclass of the
 {{#program:generate_immutable_types}}
 # This unfortunately has to be down here to prevent circular imports
 import {{program:module_path}}.thrift_metadata
-{{/program:generate_immutable_types}}
 {{> enums/thrift_enums }}
+{{/program:generate_immutable_types}}
+{{#program:generate_mutable_types}}
+from {{program:module_path}}.thrift_enums import *
+
+{{!
+    The symbol below is only required for AnyRegistry.register_module.
+    Currently, calls to that method only pass thrift_types.
+    register_module in its current form expects _fbthrift_all_enums
+    to be defined in thrift_types. For now define
+    _fbthrift_all_enums here.
+}}
+_fbthrift_all_enums = [
+    {{#program:enums}}
+    {{enum:name}},
+    {{/program:enums}}
+]
+{{/program:generate_mutable_types}}
 {{#program:generate_immutable_types}}{{!
 }}{{#program:structs}}
 

diff --git a/thrift/compiler/generate/templates/python/types/types_pyi.mustache b/thrift/compiler/generate/templates/python/types/types_pyi.mustache
@@ -47,6 +47,7 @@ import {{included_module_path}}.{{> types/types_import_path}} as {{included_modu
 {{#program:adapter_type_hint_modules}}
 import {{module_path}}
 {{/program:adapter_type_hint_modules}}
+{{#program:generate_immutable_types}}
 {{#program:enums}}
 
 class _fbthrift_compatible_with_{{enum:name}}:
@@ -69,6 +70,20 @@ class {{enum:name}}(_fbthrift_python_types.{{!
     def _to_py_deprecated(self) -> int: ...
     {{/enum:legacy_api?}}
 {{/program:enums}}
+{{/program:generate_immutable_types}}
+{{#program:generate_mutable_types}}
+{{!
+    from module_path.thrift_enums import *
+
+below does not import private symbols _fbthrift_compatible_with_*.
+These are referenced in the generated pyi files and, for simplicity,
+need to be imported through the types modules.
+}}{{#program:enums}}
+from {{program:module_path}}.thrift_enums import _fbthrift_compatible_with_{{enum:name}}
+{{/program:enums}}
+
+from {{program:module_path}}.thrift_enums import *
+{{/program:generate_mutable_types}}
 {{#program:structs}}
 
 

diff --git a/...tal_generate_abstract_types/gen-python/test/fixtures/basic/module/thrift_mutable_types.py b/...tal_generate_abstract_types/gen-python/test/fixtures/basic/module/thrift_mutable_types.py
@@ -670,67 +670,14 @@ def _to_py_deprecated(self):
             py_asyncio_types = importlib.import_module("module.ttypes")
             return thrift.util.converter.to_py_struct(py_asyncio_types.UnionToBeRenamed, self)
 
-
-
-class MyEnum(_fbthrift_python_types.Enum, int):
-    MyValue1 = 0
-    MyValue2 = 1
-    @staticmethod
-    def __get_thrift_name__() -> str:
-        return "module.MyEnum"
-
-    @staticmethod
-    def __get_thrift_uri__():
-        return "test.dev/fixtures/basic/MyEnum"
-
-    @staticmethod
-    def __get_metadata__():
-        return test.fixtures.basic.module.thrift_metadata.gen_metadata_enum_MyEnum()
-
-    def _to_python(self):
-        return self
-
-    def _to_py3(self):
-        import importlib
-        py3_types = importlib.import_module("test.fixtures.basic.module.types")
-        return py3_types.MyEnum(self.value)
-
-    def _to_py_deprecated(self):
-        return self.value
-
-
-class HackEnum(_fbthrift_python_types.Enum, int):
-    Value1 = 0
-    Value2 = 1
-    @staticmethod
-    def __get_thrift_name__() -> str:
-        return "module.HackEnum"
-
-    @staticmethod
-    def __get_thrift_uri__():
-        return "test.dev/fixtures/basic/HackEnum"
-
-    @staticmethod
-    def __get_metadata__():
-        return test.fixtures.basic.module.thrift_metadata.gen_metadata_enum_HackEnum()
-
-    def _to_python(self):
-        return self
-
-    def _to_py3(self):
-        import importlib
-        py3_types = importlib.import_module("test.fixtures.basic.module.types")
-        return py3_types.HackEnum(self.value)
-
-    def _to_py_deprecated(self):
-        return self.value
-
+from test.fixtures.basic.module.thrift_enums import *
 
 _fbthrift_all_enums = [
     MyEnum,
     HackEnum,
 ]
 
+
 _fbthrift_all_structs = [
     MyStruct,
     Containers,

diff --git a/...al_generate_abstract_types/gen-python/test/fixtures/basic/module/thrift_mutable_types.pyi b/...al_generate_abstract_types/gen-python/test/fixtures/basic/module/thrift_mutable_types.pyi
@@ -30,28 +30,10 @@ import thrift.python.types as _fbthrift_python_types
 import thrift.python.mutable_types as _fbthrift_python_mutable_types
 import thrift.python.mutable_exceptions as _fbthrift_python_mutable_exceptions
 import thrift.python.mutable_containers as _fbthrift_python_mutable_containers
+from test.fixtures.basic.module.thrift_enums import _fbthrift_compatible_with_MyEnum
+from test.fixtures.basic.module.thrift_enums import _fbthrift_compatible_with_HackEnum
 
-class _fbthrift_compatible_with_MyEnum:
-    pass
-
-
-class MyEnum(_fbthrift_python_types.Enum, int, _fbthrift_python_abstract_types.MyEnum, _fbthrift_compatible_with_MyEnum):
-    MyValue1: MyEnum = ...
-    MyValue2: MyEnum = ...
-    def _to_python(self) -> MyEnum: ...
-    def _to_py3(self) -> "test.fixtures.basic.module.types.MyEnum": ...  # type: ignore
-    def _to_py_deprecated(self) -> int: ...
-
-class _fbthrift_compatible_with_HackEnum:
-    pass
-
-
-class HackEnum(_fbthrift_python_types.Enum, int, _fbthrift_python_abstract_types.HackEnum, _fbthrift_compatible_with_HackEnum):
-    Value1: HackEnum = ...
-    Value2: HackEnum = ...
-    def _to_python(self) -> HackEnum: ...
-    def _to_py3(self) -> "test.fixtures.basic.module.types.HackEnum": ...  # type: ignore
-    def _to_py_deprecated(self) -> int: ...
+from test.fixtures.basic.module.thrift_enums import *
 
 
 class _fbthrift_compatible_with_MyStruct:

diff --git a/...ntal_generate_mutable_types/gen-python/test/fixtures/basic/module/thrift_mutable_types.py b/...ntal_generate_mutable_types/gen-python/test/fixtures/basic/module/thrift_mutable_types.py
@@ -655,67 +655,14 @@ def _to_py_deprecated(self):
             py_asyncio_types = importlib.import_module("module.ttypes")
             return thrift.util.converter.to_py_struct(py_asyncio_types.UnionToBeRenamed, self)
 
-
-
-class MyEnum(_fbthrift_python_types.Enum, int):
-    MyValue1 = 0
-    MyValue2 = 1
-    @staticmethod
-    def __get_thrift_name__() -> str:
-        return "module.MyEnum"
-
-    @staticmethod
-    def __get_thrift_uri__():
-        return "test.dev/fixtures/basic/MyEnum"
-
-    @staticmethod
-    def __get_metadata__():
-        return test.fixtures.basic.module.thrift_metadata.gen_metadata_enum_MyEnum()
-
-    def _to_python(self):
-        return self
-
-    def _to_py3(self):
-        import importlib
-        py3_types = importlib.import_module("test.fixtures.basic.module.types")
-        return py3_types.MyEnum(self.value)
-
-    def _to_py_deprecated(self):
-        return self.value
-
-
-class HackEnum(_fbthrift_python_types.Enum, int):
-    Value1 = 0
-    Value2 = 1
-    @staticmethod
-    def __get_thrift_name__() -> str:
-        return "module.HackEnum"
-
-    @staticmethod
-    def __get_thrift_uri__():
-        return "test.dev/fixtures/basic/HackEnum"
-
-    @staticmethod
-    def __get_metadata__():
-        return test.fixtures.basic.module.thrift_metadata.gen_metadata_enum_HackEnum()
-
-    def _to_python(self):
-        return self
-
-    def _to_py3(self):
-        import importlib
-        py3_types = importlib.import_module("test.fixtures.basic.module.types")
-        return py3_types.HackEnum(self.value)
-
-    def _to_py_deprecated(self):
-        return self.value
-
+from test.fixtures.basic.module.thrift_enums import *
 
 _fbthrift_all_enums = [
     MyEnum,
     HackEnum,
 ]
 
+
 _fbthrift_all_structs = [
     MyStruct,
     Containers,

diff --git a/...tal_generate_mutable_types/gen-python/test/fixtures/basic/module/thrift_mutable_types.pyi b/...tal_generate_mutable_types/gen-python/test/fixtures/basic/module/thrift_mutable_types.pyi
@@ -29,28 +29,10 @@ import thrift.python.types as _fbthrift_python_types
 import thrift.python.mutable_types as _fbthrift_python_mutable_types
 import thrift.python.mutable_exceptions as _fbthrift_python_mutable_exceptions
 import thrift.python.mutable_containers as _fbthrift_python_mutable_containers
+from test.fixtures.basic.module.thrift_enums import _fbthrift_compatible_with_MyEnum
+from test.fixtures.basic.module.thrift_enums import _fbthrift_compatible_with_HackEnum
 
-class _fbthrift_compatible_with_MyEnum:
-    pass
-
-
-class MyEnum(_fbthrift_python_types.Enum, int, _fbthrift_compatible_with_MyEnum):
-    MyValue1: MyEnum = ...
-    MyValue2: MyEnum = ...
-    def _to_python(self) -> MyEnum: ...
-    def _to_py3(self) -> "test.fixtures.basic.module.types.MyEnum": ...  # type: ignore
-    def _to_py_deprecated(self) -> int: ...
-
-class _fbthrift_compatible_with_HackEnum:
-    pass
-
-
-class HackEnum(_fbthrift_python_types.Enum, int, _fbthrift_compatible_with_HackEnum):
-    Value1: HackEnum = ...
-    Value2: HackEnum = ...
-    def _to_python(self) -> HackEnum: ...
-    def _to_py3(self) -> "test.fixtures.basic.module.types.HackEnum": ...  # type: ignore
-    def _to_py_deprecated(self) -> int: ...
+from test.fixtures.basic.module.thrift_enums import *
 
 
 class _fbthrift_compatible_with_MyStruct:

diff --git a/...rate_mutable_types_service/gen-python/meta/example/thrift/service/thrift_mutable_types.py b/...rate_mutable_types_service/gen-python/meta/example/thrift/service/thrift_mutable_types.py
@@ -180,11 +180,12 @@ def _to_py_deprecated(self):
             py_asyncio_types = importlib.import_module("service.ttypes")
             return thrift.util.converter.to_py_struct(py_asyncio_types.WhisperException, self)
 
-
+from meta.example.thrift.service.thrift_enums import *
 
 _fbthrift_all_enums = [
 ]
 
+
 _fbthrift_all_structs = [
     EchoRequest,
     EchoResponse,

diff --git a/...ate_mutable_types_service/gen-python/meta/example/thrift/service/thrift_mutable_types.pyi b/...ate_mutable_types_service/gen-python/meta/example/thrift/service/thrift_mutable_types.pyi
@@ -28,6 +28,8 @@ import thrift.python.mutable_types as _fbthrift_python_mutable_types
 import thrift.python.mutable_exceptions as _fbthrift_python_mutable_exceptions
 import thrift.python.mutable_containers as _fbthrift_python_mutable_containers
 
+from meta.example.thrift.service.thrift_enums import *
+
 
 class _fbthrift_compatible_with_EchoRequest:
     pass

diff --git a/...out/python_experimental_generate_abstract_types/gen-python/module/thrift_mutable_types.py b/...out/python_experimental_generate_abstract_types/gen-python/module/thrift_mutable_types.py
@@ -1884,42 +1884,13 @@ def _to_py_deprecated(self):
             py_asyncio_types = importlib.import_module("module.ttypes")
             return thrift.util.converter.to_py_struct(py_asyncio_types.complexException, self)
 
-
-
-class MyEnum(_fbthrift_python_types.Enum, int):
-    MyValue1 = 0
-    MyValue2 = 1
-    MyValue3 = 3
-    MyValue4 = 4
-    MyValue5 = 5
-    @staticmethod
-    def __get_thrift_name__() -> str:
-        return "module.MyEnum"
-
-    @staticmethod
-    def __get_thrift_uri__():
-        return None
-
-    @staticmethod
-    def __get_metadata__():
-        return module.thrift_metadata.gen_metadata_enum_MyEnum()
-
-    def _to_python(self):
-        return self
-
-    def _to_py3(self):
-        import importlib
-        py3_types = importlib.import_module("module.types")
-        return py3_types.MyEnum(self.value)
-
-    def _to_py_deprecated(self):
-        return self.value
-
+from module.thrift_enums import *
 
 _fbthrift_all_enums = [
     MyEnum,
 ]
 
+
 _fbthrift_all_structs = [
     MyStructFloatFieldThrowExp,
     MyStructMapFloatThrowExp,

diff --git a/...ut/python_experimental_generate_abstract_types/gen-python/module/thrift_mutable_types.pyi b/...ut/python_experimental_generate_abstract_types/gen-python/module/thrift_mutable_types.pyi
@@ -30,20 +30,9 @@ import thrift.python.types as _fbthrift_python_types
 import thrift.python.mutable_types as _fbthrift_python_mutable_types
 import thrift.python.mutable_exceptions as _fbthrift_python_mutable_exceptions
 import thrift.python.mutable_containers as _fbthrift_python_mutable_containers
+from module.thrift_enums import _fbthrift_compatible_with_MyEnum
 
-class _fbthrift_compatible_with_MyEnum:
-    pass
-
-
-class MyEnum(_fbthrift_python_types.Enum, int, _fbthrift_python_abstract_types.MyEnum, _fbthrift_compatible_with_MyEnum):
-    MyValue1: MyEnum = ...
-    MyValue2: MyEnum = ...
-    MyValue3: MyEnum = ...
-    MyValue4: MyEnum = ...
-    MyValue5: MyEnum = ...
-    def _to_python(self) -> MyEnum: ...
-    def _to_py3(self) -> "module.types.MyEnum": ...  # type: ignore
-    def _to_py_deprecated(self) -> int: ...
+from module.thrift_enums import *
 
 
 class _fbthrift_compatible_with_MyStructFloatFieldThrowExp:

diff --git a/.../out/python_experimental_generate_mutable_types/gen-python/module/thrift_mutable_types.py b/.../out/python_experimental_generate_mutable_types/gen-python/module/thrift_mutable_types.py
@@ -1862,42 +1862,13 @@ def _to_py_deprecated(self):
             py_asyncio_types = importlib.import_module("module.ttypes")
             return thrift.util.converter.to_py_struct(py_asyncio_types.complexException, self)
 
-
-
-class MyEnum(_fbthrift_python_types.Enum, int):
-    MyValue1 = 0
-    MyValue2 = 1
-    MyValue3 = 3
-    MyValue4 = 4
-    MyValue5 = 5
-    @staticmethod
-    def __get_thrift_name__() -> str:
-        return "module.MyEnum"
-
-    @staticmethod
-    def __get_thrift_uri__():
-        return None
-
-    @staticmethod
-    def __get_metadata__():
-        return module.thrift_metadata.gen_metadata_enum_MyEnum()
-
-    def _to_python(self):
-        return self
-
-    def _to_py3(self):
-        import importlib
-        py3_types = importlib.import_module("module.types")
-        return py3_types.MyEnum(self.value)
-
-    def _to_py_deprecated(self):
-        return self.value
-
+from module.thrift_enums import *
 
 _fbthrift_all_enums = [
     MyEnum,
 ]
 
+
 _fbthrift_all_structs = [
     MyStructFloatFieldThrowExp,
     MyStructMapFloatThrowExp,