Implement union specific enum and properties in abstract unions

Summary: See title. Reviewed By: yoney Differential Revision: D65348204 fbshipit-source-id: f2a06e990fc80cfaa60af42231b1ca896ef65942
facebook · Nov 2, 2024 · 9a167d2 · 9a167d2
1 parent 18aaaeb
commit 9a167d2
Show file tree

Hide file tree

Showing 54 changed files with 617 additions and 441 deletions.
diff --git a/...ty/thrift/src/thrift/compiler/generate/templates/python/thrift_abstract_types.py.mustache b/...ty/thrift/src/thrift/compiler/generate/templates/python/thrift_abstract_types.py.mustache
@@ -25,6 +25,10 @@ Generated Python classes for abstract thrift-python types.
 
 import abc as _abc
 import typing as _typing
+{{#program:unions?}}
+import enum as _enum
+
+{{/program:unions?}}
 
 import folly.iobuf as _fbthrift_iobuf
 {{#program:include_namespaces}}
@@ -82,6 +86,23 @@ class {{struct:py_name}}({{!
     @_abc.abstractmethod
     def _to_py_deprecated(self) -> "{{program:py_deprecated_module_path}}.ttypes.{{struct:py_name}}": ...  # type: ignore
     {{/struct:legacy_api?}}
+    {{#struct:union?}}
+
+    class FbThriftUnionFieldEnum(_enum.Enum):
+        EMPTY: {{struct:py_name}}.FbThriftUnionFieldEnum = 0
+    {{#struct:fields_ordered_by_id}}
+        {{field:py_name}}: {{struct:py_name}}.FbThriftUnionFieldEnum = {{field:key}}
+    {{/struct:fields_ordered_by_id}}
+
+    {{! For backwards compatibility, set the name to the name
+        used in immutable thrift-python, which is the name
+        of the Union (sigh!)
+    }}FbThriftUnionFieldEnum.__name__ = "{{struct:py_name}}"
+
+    fbthrift_current_value: _typing.Final[{{ > types/field_value_pep484_union_type}}]
+    fbthrift_current_field: _typing.Final[FbThriftUnionFieldEnum]
+
+    {{/struct:union?}}
 {{#struct:has_adapter?}}
 # TODO: Figure out how to handle adapters
 # {{struct:py_name}} = {{adapter:type_hint}}{{#adapter:is_generic?}}[{{> structs/unadapted_name}}]{{/adapter:is_generic?}}

diff --git a/...rc/thrift/compiler/generate/templates/python/types/field_value_pep484_union_type.mustache b/...rc/thrift/compiler/generate/templates/python/types/field_value_pep484_union_type.mustache
@@ -17,5 +17,12 @@
 }}{{!
 Typing for a union field
 }}_typing.Union[None{{#struct:fields_ordered_by_id}}{{!
-  }}, {{#field:type}}{{> types/pep484_type}}{{/field:type}}{{!
+  }}, {{#field:type}}{{!
+    Do not use adapted abstract types until abstract types support them.
+  }}{{#program:generate_abstract_types}}{{!
+  }}{{> types/unadapted_pep484_type}}{{!
+  }}{{/program:generate_abstract_types}}{{!
+  }}{{^program:generate_abstract_types}}{{!
+  }}{{> types/pep484_type}}{{!
+  }}{{/program:generate_abstract_types}}{{/field:type}}{{!
   }}{{/struct:fields_ordered_by_id}}]
diff --git a/third-party/thrift/src/thrift/compiler/generate/templates/python/types/types.mustache b/third-party/thrift/src/thrift/compiler/generate/templates/python/types/types.mustache
@@ -21,7 +21,7 @@ Generated Python classes for Thrift types
 
 import folly.iobuf as _fbthrift_iobuf
 {{#program:generate_unified_thrift_python_type_hints?}}
-import {{program:module_path}}.thrift_abstract_types
+import {{program:module_path}}.thrift_abstract_types as _fbthrift_abstract_types
 {{/program:generate_unified_thrift_python_type_hints?}}
 import {{program:base_library_package}}.types as _fbthrift_python_types
 {{#program:generate_immutable_types}}
@@ -52,7 +52,7 @@ If the class uses the abstract class as a direct baseclass, it results in the er
 
 Until we resolve the above error, register the class as virtual subclass of the abstract class.
 }}
-@{{program:module_path}}.thrift_abstract_types.{{struct:py_name}}.register
+@_fbthrift_abstract_types.{{struct:py_name}}.register
 {{/program:generate_unified_thrift_python_type_hints?}}{{!
 }}class {{> structs/unadapted_name}}{{!
 }}(metaclass={{!
@@ -77,7 +77,16 @@ Until we resolve the above error, register the class as virtual subclass of the
             {{field:idl_type}}, # IDL type (see BaseTypeEnum)
         ),
         {{/struct:fields_ordered_by_id}}
-    )
+    ){{!
+    }}{{#struct:union?}}{{!
+    }}{{#program:generate_unified_thrift_python_type_hints?}}{{!
+    }}
+
+    _fbthrift_union_field_enum = _fbthrift_abstract_types.{{struct:py_name}}.FbThriftUnionFieldEnum
+    {{!
+    }}{{/program:generate_unified_thrift_python_type_hints?}}{{!
+    }}{{/struct:union?}}
+
 
     @staticmethod
     def __get_thrift_name__() -> str:

diff --git a/third-party/thrift/src/thrift/compiler/generate/templates/python/types/types_pyi.mustache b/third-party/thrift/src/thrift/compiler/generate/templates/python/types/types_pyi.mustache
@@ -120,6 +120,7 @@ class {{> structs/unadapted_name}}({{!
     {{/struct:plain?}}
     {{#struct:union?}}
 
+{{^program:generate_unified_thrift_python_type_hints?}}
 {{#program:generate_immutable_types}}
     class Type(enum.Enum):
         EMPTY: {{> structs/unadapted_name}}.Type = ...
@@ -138,6 +139,12 @@ class {{> structs/unadapted_name}}({{!
 {{#program:generate_immutable_types}}
     FbThriftUnionFieldEnum = Type
 {{/program:generate_immutable_types}}
+{{/program:generate_unified_thrift_python_type_hints?}}
+{{#program:generate_unified_thrift_python_type_hints?}}
+{{#program:generate_immutable_types}}
+    Type = FbThriftUnionFieldEnum
+{{/program:generate_immutable_types}}
+{{/program:generate_unified_thrift_python_type_hints?}}
 
 {{#program:generate_immutable_types}}
     @classmethod

diff --git a/...y_thrift_python_type_hints/gen-python/test/fixtures/basic/module/thrift_abstract_types.py b/...y_thrift_python_type_hints/gen-python/test/fixtures/basic/module/thrift_abstract_types.py
@@ -13,6 +13,8 @@
 
 import abc as _abc
 import typing as _typing
+import enum as _enum
+
 
 import folly.iobuf as _fbthrift_iobuf
 
@@ -118,6 +120,19 @@ def _to_py3(self) -> "test.fixtures.basic.module.types.MyUnion": ...  # type: ig
     @_abc.abstractmethod
     def _to_py_deprecated(self) -> "module.ttypes.MyUnion": ...  # type: ignore
 
+    class FbThriftUnionFieldEnum(_enum.Enum):
+        EMPTY: MyUnion.FbThriftUnionFieldEnum = 0
+        myEnum: MyUnion.FbThriftUnionFieldEnum = 1
+        myStruct: MyUnion.FbThriftUnionFieldEnum = 2
+        myDataItem: MyUnion.FbThriftUnionFieldEnum = 3
+        floatSet: MyUnion.FbThriftUnionFieldEnum = 4
+
+    FbThriftUnionFieldEnum.__name__ = "MyUnion"
+
+    fbthrift_current_value: _typing.Final[_typing.Union[None, MyEnum, MyStruct, MyDataItem, _typing.AbstractSet[float]]]
+    fbthrift_current_field: _typing.Final[FbThriftUnionFieldEnum]
+
+
 class MyException(_abc.ABC):
     @property
     @_abc.abstractmethod
@@ -194,5 +209,15 @@ def _to_py3(self) -> "test.fixtures.basic.module.types.UnionToBeRenamed": ...  #
     @_abc.abstractmethod
     def _to_py_deprecated(self) -> "module.ttypes.UnionToBeRenamed": ...  # type: ignore
 
+    class FbThriftUnionFieldEnum(_enum.Enum):
+        EMPTY: UnionToBeRenamed.FbThriftUnionFieldEnum = 0
+        reserved_field: UnionToBeRenamed.FbThriftUnionFieldEnum = 1
+
+    FbThriftUnionFieldEnum.__name__ = "UnionToBeRenamed"
+
+    fbthrift_current_value: _typing.Final[_typing.Union[None, int]]
+    fbthrift_current_field: _typing.Final[FbThriftUnionFieldEnum]
+
+
 MyEnumAlias = MyEnum
 MyDataItemAlias = MyDataItem
diff --git a/...fy_thrift_python_type_hints/gen-python/test/fixtures/basic/module/thrift_mutable_types.py b/...fy_thrift_python_type_hints/gen-python/test/fixtures/basic/module/thrift_mutable_types.py
@@ -12,15 +12,15 @@
 from __future__ import annotations
 
 import folly.iobuf as _fbthrift_iobuf
-import test.fixtures.basic.module.thrift_abstract_types
+import test.fixtures.basic.module.thrift_abstract_types as _fbthrift_abstract_types
 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_typeinfos as _fbthrift_python_mutable_typeinfos
 
 
 
-@test.fixtures.basic.module.thrift_abstract_types.MyStruct.register
+@_fbthrift_abstract_types.MyStruct.register
 class MyStruct(metaclass=_fbthrift_python_mutable_types.MutableStructMeta):
     _fbthrift_SPEC = (
         _fbthrift_python_types.FieldInfo(
@@ -157,7 +157,7 @@ def _to_py_deprecated(self):
             return thrift.util.converter.to_py_struct(py_asyncio_types.MyStruct, self)
 
 
-@test.fixtures.basic.module.thrift_abstract_types.Containers.register
+@_fbthrift_abstract_types.Containers.register
 class Containers(metaclass=_fbthrift_python_mutable_types.MutableStructMeta):
     _fbthrift_SPEC = (
         _fbthrift_python_types.FieldInfo(
@@ -228,7 +228,7 @@ def _to_py_deprecated(self):
             return thrift.util.converter.to_py_struct(py_asyncio_types.Containers, self)
 
 
-@test.fixtures.basic.module.thrift_abstract_types.MyDataItem.register
+@_fbthrift_abstract_types.MyDataItem.register
 class MyDataItem(metaclass=_fbthrift_python_mutable_types.MutableStructMeta):
     _fbthrift_SPEC = (
     )
@@ -266,7 +266,7 @@ def _to_py_deprecated(self):
             return thrift.util.converter.to_py_struct(py_asyncio_types.MyDataItem, self)
 
 
-@test.fixtures.basic.module.thrift_abstract_types.MyUnion.register
+@_fbthrift_abstract_types.MyUnion.register
 class MyUnion(metaclass=_fbthrift_python_mutable_types.MutableUnionMeta):
     _fbthrift_SPEC = (
         _fbthrift_python_types.FieldInfo(
@@ -315,6 +315,9 @@ class MyUnion(metaclass=_fbthrift_python_mutable_types.MutableUnionMeta):
         ),
     )
 
+    _fbthrift_union_field_enum = _fbthrift_abstract_types.MyUnion.FbThriftUnionFieldEnum
+
+
     @staticmethod
     def __get_thrift_name__() -> str:
         return "module.MyUnion"
@@ -348,7 +351,7 @@ def _to_py_deprecated(self):
             return thrift.util.converter.to_py_struct(py_asyncio_types.MyUnion, self)
 
 
-@test.fixtures.basic.module.thrift_abstract_types.MyException.register
+@_fbthrift_abstract_types.MyException.register
 class MyException(metaclass=_fbthrift_python_mutable_exceptions.MutableGeneratedErrorMeta):
     _fbthrift_SPEC = (
         _fbthrift_python_types.FieldInfo(
@@ -430,7 +433,7 @@ def _to_py_deprecated(self):
             return thrift.util.converter.to_py_struct(py_asyncio_types.MyException, self)
 
 
-@test.fixtures.basic.module.thrift_abstract_types.MyExceptionWithMessage.register
+@_fbthrift_abstract_types.MyExceptionWithMessage.register
 class MyExceptionWithMessage(metaclass=_fbthrift_python_mutable_exceptions.MutableGeneratedErrorMeta):
     _fbthrift_SPEC = (
         _fbthrift_python_types.FieldInfo(
@@ -519,7 +522,7 @@ def _to_py_deprecated(self):
             return thrift.util.converter.to_py_struct(py_asyncio_types.MyExceptionWithMessage, self)
 
 
-@test.fixtures.basic.module.thrift_abstract_types.ReservedKeyword.register
+@_fbthrift_abstract_types.ReservedKeyword.register
 class ReservedKeyword(metaclass=_fbthrift_python_mutable_types.MutableStructMeta):
     _fbthrift_SPEC = (
         _fbthrift_python_types.FieldInfo(
@@ -568,7 +571,7 @@ def _to_py_deprecated(self):
             return thrift.util.converter.to_py_struct(py_asyncio_types.ReservedKeyword, self)
 
 
-@test.fixtures.basic.module.thrift_abstract_types.UnionToBeRenamed.register
+@_fbthrift_abstract_types.UnionToBeRenamed.register
 class UnionToBeRenamed(metaclass=_fbthrift_python_mutable_types.MutableUnionMeta):
     _fbthrift_SPEC = (
         _fbthrift_python_types.FieldInfo(
@@ -584,6 +587,9 @@ class UnionToBeRenamed(metaclass=_fbthrift_python_mutable_types.MutableUnionMeta
         ),
     )
 
+    _fbthrift_union_field_enum = _fbthrift_abstract_types.UnionToBeRenamed.FbThriftUnionFieldEnum
+
+
     @staticmethod
     def __get_thrift_name__() -> str:
         return "module.UnionToBeRenamed"

diff --git a/...y_thrift_python_type_hints/gen-python/test/fixtures/basic/module/thrift_mutable_types.pyi b/...y_thrift_python_type_hints/gen-python/test/fixtures/basic/module/thrift_mutable_types.pyi
@@ -190,13 +190,6 @@ class MyUnion(_fbthrift_python_mutable_types.MutableUnion, _fbthrift_compatible_
     ) -> None: ...
 
 
-    class FbThriftUnionFieldEnum(enum.Enum):
-        EMPTY: MyUnion.FbThriftUnionFieldEnum = ...
-        myEnum: MyUnion.FbThriftUnionFieldEnum = ...
-        myStruct: MyUnion.FbThriftUnionFieldEnum = ...
-        myDataItem: MyUnion.FbThriftUnionFieldEnum = ...
-        floatSet: MyUnion.FbThriftUnionFieldEnum = ...
-
 
     fbthrift_current_value: _typing.Final[_typing.Union[None, MyEnum, MyStruct, MyDataItem, _fbthrift_python_mutable_containers.MutableSet[float]]]
     fbthrift_current_field: _typing.Final[FbThriftUnionFieldEnum]
@@ -289,10 +282,6 @@ class UnionToBeRenamed(_fbthrift_python_mutable_types.MutableUnion, _fbthrift_co
     ) -> None: ...
 
 
-    class FbThriftUnionFieldEnum(enum.Enum):
-        EMPTY: UnionToBeRenamed.FbThriftUnionFieldEnum = ...
-        reserved_field: UnionToBeRenamed.FbThriftUnionFieldEnum = ...
-
 
     fbthrift_current_value: _typing.Final[_typing.Union[None, int]]
     fbthrift_current_field: _typing.Final[FbThriftUnionFieldEnum]

diff --git a/...ntal_unify_thrift_python_type_hints/gen-python/test/fixtures/basic/module/thrift_types.py b/...ntal_unify_thrift_python_type_hints/gen-python/test/fixtures/basic/module/thrift_types.py
@@ -8,13 +8,13 @@
 from __future__ import annotations
 
 import folly.iobuf as _fbthrift_iobuf
-import test.fixtures.basic.module.thrift_abstract_types
+import test.fixtures.basic.module.thrift_abstract_types as _fbthrift_abstract_types
 import thrift.python.types as _fbthrift_python_types
 import thrift.python.exceptions as _fbthrift_python_exceptions
 
 
 
-@test.fixtures.basic.module.thrift_abstract_types.MyStruct.register
+@_fbthrift_abstract_types.MyStruct.register
 class MyStruct(metaclass=_fbthrift_python_types.StructMeta):
     _fbthrift_SPEC = (
         _fbthrift_python_types.FieldInfo(
@@ -156,7 +156,7 @@ def _to_py_deprecated(self):
             return thrift.util.converter.to_py_struct(py_asyncio_types.MyStruct, self)
 
 
-@test.fixtures.basic.module.thrift_abstract_types.Containers.register
+@_fbthrift_abstract_types.Containers.register
 class Containers(metaclass=_fbthrift_python_types.StructMeta):
     _fbthrift_SPEC = (
         _fbthrift_python_types.FieldInfo(
@@ -232,7 +232,7 @@ def _to_py_deprecated(self):
             return thrift.util.converter.to_py_struct(py_asyncio_types.Containers, self)
 
 
-@test.fixtures.basic.module.thrift_abstract_types.MyDataItem.register
+@_fbthrift_abstract_types.MyDataItem.register
 class MyDataItem(metaclass=_fbthrift_python_types.StructMeta):
     _fbthrift_SPEC = (
     )
@@ -275,7 +275,7 @@ def _to_py_deprecated(self):
             return thrift.util.converter.to_py_struct(py_asyncio_types.MyDataItem, self)
 
 
-@test.fixtures.basic.module.thrift_abstract_types.MyUnion.register
+@_fbthrift_abstract_types.MyUnion.register
 class MyUnion(metaclass=_fbthrift_python_types.UnionMeta):
     _fbthrift_SPEC = (
         _fbthrift_python_types.FieldInfo(
@@ -324,6 +324,9 @@ class MyUnion(metaclass=_fbthrift_python_types.UnionMeta):
         ),
     )
 
+    _fbthrift_union_field_enum = _fbthrift_abstract_types.MyUnion.FbThriftUnionFieldEnum
+
+
     @staticmethod
     def __get_thrift_name__() -> str:
         return "module.MyUnion"
@@ -362,7 +365,7 @@ def _to_py_deprecated(self):
             return thrift.util.converter.to_py_struct(py_asyncio_types.MyUnion, self)
 
 
-@test.fixtures.basic.module.thrift_abstract_types.MyException.register
+@_fbthrift_abstract_types.MyException.register
 class MyException(metaclass=_fbthrift_python_exceptions.GeneratedErrorMeta):
     _fbthrift_SPEC = (
         _fbthrift_python_types.FieldInfo(
@@ -449,7 +452,7 @@ def _to_py_deprecated(self):
             return thrift.util.converter.to_py_struct(py_asyncio_types.MyException, self)
 
 
-@test.fixtures.basic.module.thrift_abstract_types.MyExceptionWithMessage.register
+@_fbthrift_abstract_types.MyExceptionWithMessage.register
 class MyExceptionWithMessage(metaclass=_fbthrift_python_exceptions.GeneratedErrorMeta):
     _fbthrift_SPEC = (
         _fbthrift_python_types.FieldInfo(
@@ -543,7 +546,7 @@ def _to_py_deprecated(self):
             return thrift.util.converter.to_py_struct(py_asyncio_types.MyExceptionWithMessage, self)
 
 
-@test.fixtures.basic.module.thrift_abstract_types.ReservedKeyword.register
+@_fbthrift_abstract_types.ReservedKeyword.register
 class ReservedKeyword(metaclass=_fbthrift_python_types.StructMeta):
     _fbthrift_SPEC = (
         _fbthrift_python_types.FieldInfo(
@@ -597,7 +600,7 @@ def _to_py_deprecated(self):
             return thrift.util.converter.to_py_struct(py_asyncio_types.ReservedKeyword, self)
 
 
-@test.fixtures.basic.module.thrift_abstract_types.UnionToBeRenamed.register
+@_fbthrift_abstract_types.UnionToBeRenamed.register
 class UnionToBeRenamed(metaclass=_fbthrift_python_types.UnionMeta):
     _fbthrift_SPEC = (
         _fbthrift_python_types.FieldInfo(
@@ -613,6 +616,9 @@ class UnionToBeRenamed(metaclass=_fbthrift_python_types.UnionMeta):
         ),
     )
 
+    _fbthrift_union_field_enum = _fbthrift_abstract_types.UnionToBeRenamed.FbThriftUnionFieldEnum
+
+
     @staticmethod
     def __get_thrift_name__() -> str:
         return "module.UnionToBeRenamed"