Support UnionPatch in python

Summary: Previously unions were treated as structs. This diff added the support with unit-test. Reviewed By: thedavekwon Differential Revision: D66324299 fbshipit-source-id: 704ed09e2782c968082898efbfa3393b1dc117af
facebook · Nov 25, 2024 · eca7bf3 · eca7bf3
1 parent e0b4f8d
commit eca7bf3
Show file tree

Hide file tree

Showing 2 changed files with 61 additions and 30 deletions.
diff --git a/thrift/compiler/generate/templates/patch/thrift_patch.py.mustache b/thrift/compiler/generate/templates/patch/thrift_patch.py.mustache
@@ -23,6 +23,7 @@ import typing as _typing
 
 from common.thrift.patch.detail.dynamic_patch import (
     BaseStructPatch,
+    BaseUnionPatch,
     ListPatch,
     SetPatch,
     MapPatch,
@@ -40,6 +41,7 @@ from common.thrift.patch.detail.py_bindings.DynamicPatch import (
     StringPatch,
     BinaryPatch,
     StructPatch as DynamicStructPatch,
+    UnionPatch as DynamicUnionPatch,
 )
 
 import thrift.python.types as _fbthrift_python_types
@@ -59,7 +61,7 @@ import {{included_module_path}}.thrift_patch
 {{#program:structs}}
 
 class {{struct:py_name}}Patch(
-    BaseStructPatch[{{program:module_mangle}}__thrift_types.{{struct:py_name}}]
+    Base{{#if struct:union?}}Union{{else}}Struct{{/if}}Patch[{{program:module_mangle}}__thrift_types.{{struct:py_name}}]
 ):
     pass
 
@@ -83,7 +85,8 @@ class {{struct:py_name}}Patch(
     @staticmethod
     def from_safe_patch(safe_patch: _fbthrift_safe_patch_types.{{struct:name}}SafePatch) -> {{struct:py_name}}Patch:
         patch = {{struct:py_name}}Patch()
-        patch._patch = DynamicStructPatch.deserialize_from_compact_protocol(safe_patch.data)
+        DynamicPatch = Dynamic{{#if struct:union?}}Union{{else}}Struct{{/if}}Patch
+        patch._patch = DynamicPatch.deserialize_from_compact_protocol(safe_patch.data)
         return patch
 
 

diff --git a/...xtures/basic/out/python_patch/gen-python-patch/test/fixtures/basic/module/thrift_patch.py b/...xtures/basic/out/python_patch/gen-python-patch/test/fixtures/basic/module/thrift_patch.py
@@ -6,6 +6,7 @@
 
 from common.thrift.patch.detail.dynamic_patch import (
     BaseStructPatch,
+    BaseUnionPatch,
     ListPatch,
     SetPatch,
     MapPatch,
@@ -23,6 +24,7 @@
     StringPatch,
     BinaryPatch,
     StructPatch as DynamicStructPatch,
+    UnionPatch as DynamicUnionPatch,
 )
 
 import thrift.python.types as _fbthrift_python_types
@@ -57,7 +59,8 @@ def to_safe_patch(self) -> _fbthrift_safe_patch_types.MyStructSafePatch:
     @staticmethod
     def from_safe_patch(safe_patch: _fbthrift_safe_patch_types.MyStructSafePatch) -> MyStructPatch:
         patch = MyStructPatch()
-        patch._patch = DynamicStructPatch.deserialize_from_compact_protocol(safe_patch.data)
+        DynamicPatch = DynamicStructPatch
+        patch._patch = DynamicPatch.deserialize_from_compact_protocol(safe_patch.data)
         return patch
 
 
@@ -80,7 +83,8 @@ def to_safe_patch(self) -> _fbthrift_safe_patch_types.MyStructSafePatch:
     @staticmethod
     def from_safe_patch(safe_patch: _fbthrift_safe_patch_types.MyStructSafePatch) -> MyStructPatch:
         patch = MyStructPatch()
-        patch._patch = DynamicStructPatch.deserialize_from_compact_protocol(safe_patch.data)
+        DynamicPatch = DynamicStructPatch
+        patch._patch = DynamicPatch.deserialize_from_compact_protocol(safe_patch.data)
         return patch
 
 
@@ -103,7 +107,8 @@ def to_safe_patch(self) -> _fbthrift_safe_patch_types.MyStructSafePatch:
     @staticmethod
     def from_safe_patch(safe_patch: _fbthrift_safe_patch_types.MyStructSafePatch) -> MyStructPatch:
         patch = MyStructPatch()
-        patch._patch = DynamicStructPatch.deserialize_from_compact_protocol(safe_patch.data)
+        DynamicPatch = DynamicStructPatch
+        patch._patch = DynamicPatch.deserialize_from_compact_protocol(safe_patch.data)
         return patch
 
 
@@ -126,7 +131,8 @@ def to_safe_patch(self) -> _fbthrift_safe_patch_types.MyStructSafePatch:
     @staticmethod
     def from_safe_patch(safe_patch: _fbthrift_safe_patch_types.MyStructSafePatch) -> MyStructPatch:
         patch = MyStructPatch()
-        patch._patch = DynamicStructPatch.deserialize_from_compact_protocol(safe_patch.data)
+        DynamicPatch = DynamicStructPatch
+        patch._patch = DynamicPatch.deserialize_from_compact_protocol(safe_patch.data)
         return patch
 
 
@@ -149,7 +155,8 @@ def to_safe_patch(self) -> _fbthrift_safe_patch_types.MyStructSafePatch:
     @staticmethod
     def from_safe_patch(safe_patch: _fbthrift_safe_patch_types.MyStructSafePatch) -> MyStructPatch:
         patch = MyStructPatch()
-        patch._patch = DynamicStructPatch.deserialize_from_compact_protocol(safe_patch.data)
+        DynamicPatch = DynamicStructPatch
+        patch._patch = DynamicPatch.deserialize_from_compact_protocol(safe_patch.data)
         return patch
 
 
@@ -172,7 +179,8 @@ def to_safe_patch(self) -> _fbthrift_safe_patch_types.MyStructSafePatch:
     @staticmethod
     def from_safe_patch(safe_patch: _fbthrift_safe_patch_types.MyStructSafePatch) -> MyStructPatch:
         patch = MyStructPatch()
-        patch._patch = DynamicStructPatch.deserialize_from_compact_protocol(safe_patch.data)
+        DynamicPatch = DynamicStructPatch
+        patch._patch = DynamicPatch.deserialize_from_compact_protocol(safe_patch.data)
         return patch
 
 
@@ -195,7 +203,8 @@ def to_safe_patch(self) -> _fbthrift_safe_patch_types.MyStructSafePatch:
     @staticmethod
     def from_safe_patch(safe_patch: _fbthrift_safe_patch_types.MyStructSafePatch) -> MyStructPatch:
         patch = MyStructPatch()
-        patch._patch = DynamicStructPatch.deserialize_from_compact_protocol(safe_patch.data)
+        DynamicPatch = DynamicStructPatch
+        patch._patch = DynamicPatch.deserialize_from_compact_protocol(safe_patch.data)
         return patch
 
 
@@ -218,7 +227,8 @@ def to_safe_patch(self) -> _fbthrift_safe_patch_types.MyStructSafePatch:
     @staticmethod
     def from_safe_patch(safe_patch: _fbthrift_safe_patch_types.MyStructSafePatch) -> MyStructPatch:
         patch = MyStructPatch()
-        patch._patch = DynamicStructPatch.deserialize_from_compact_protocol(safe_patch.data)
+        DynamicPatch = DynamicStructPatch
+        patch._patch = DynamicPatch.deserialize_from_compact_protocol(safe_patch.data)
         return patch
 
 
@@ -241,7 +251,8 @@ def to_safe_patch(self) -> _fbthrift_safe_patch_types.MyStructSafePatch:
     @staticmethod
     def from_safe_patch(safe_patch: _fbthrift_safe_patch_types.MyStructSafePatch) -> MyStructPatch:
         patch = MyStructPatch()
-        patch._patch = DynamicStructPatch.deserialize_from_compact_protocol(safe_patch.data)
+        DynamicPatch = DynamicStructPatch
+        patch._patch = DynamicPatch.deserialize_from_compact_protocol(safe_patch.data)
         return patch
 
 
@@ -270,7 +281,8 @@ def to_safe_patch(self) -> _fbthrift_safe_patch_types.ContainersSafePatch:
     @staticmethod
     def from_safe_patch(safe_patch: _fbthrift_safe_patch_types.ContainersSafePatch) -> ContainersPatch:
         patch = ContainersPatch()
-        patch._patch = DynamicStructPatch.deserialize_from_compact_protocol(safe_patch.data)
+        DynamicPatch = DynamicStructPatch
+        patch._patch = DynamicPatch.deserialize_from_compact_protocol(safe_patch.data)
         return patch
 
 
@@ -293,7 +305,8 @@ def to_safe_patch(self) -> _fbthrift_safe_patch_types.ContainersSafePatch:
     @staticmethod
     def from_safe_patch(safe_patch: _fbthrift_safe_patch_types.ContainersSafePatch) -> ContainersPatch:
         patch = ContainersPatch()
-        patch._patch = DynamicStructPatch.deserialize_from_compact_protocol(safe_patch.data)
+        DynamicPatch = DynamicStructPatch
+        patch._patch = DynamicPatch.deserialize_from_compact_protocol(safe_patch.data)
         return patch
 
 
@@ -316,7 +329,8 @@ def to_safe_patch(self) -> _fbthrift_safe_patch_types.ContainersSafePatch:
     @staticmethod
     def from_safe_patch(safe_patch: _fbthrift_safe_patch_types.ContainersSafePatch) -> ContainersPatch:
         patch = ContainersPatch()
-        patch._patch = DynamicStructPatch.deserialize_from_compact_protocol(safe_patch.data)
+        DynamicPatch = DynamicStructPatch
+        patch._patch = DynamicPatch.deserialize_from_compact_protocol(safe_patch.data)
         return patch
 
 
@@ -328,7 +342,7 @@ class MyDataItemPatch(
 
 
 class MyUnionPatch(
-    BaseStructPatch[_fbthrift__test__fixtures__basic__module__thrift_types.MyUnion]
+    BaseUnionPatch[_fbthrift__test__fixtures__basic__module__thrift_types.MyUnion]
 ):
     pass
 
@@ -351,7 +365,8 @@ def to_safe_patch(self) -> _fbthrift_safe_patch_types.MyUnionSafePatch:
     @staticmethod
     def from_safe_patch(safe_patch: _fbthrift_safe_patch_types.MyUnionSafePatch) -> MyUnionPatch:
         patch = MyUnionPatch()
-        patch._patch = DynamicStructPatch.deserialize_from_compact_protocol(safe_patch.data)
+        DynamicPatch = DynamicUnionPatch
+        patch._patch = DynamicPatch.deserialize_from_compact_protocol(safe_patch.data)
         return patch
 
 
@@ -374,7 +389,8 @@ def to_safe_patch(self) -> _fbthrift_safe_patch_types.MyUnionSafePatch:
     @staticmethod
     def from_safe_patch(safe_patch: _fbthrift_safe_patch_types.MyUnionSafePatch) -> MyUnionPatch:
         patch = MyUnionPatch()
-        patch._patch = DynamicStructPatch.deserialize_from_compact_protocol(safe_patch.data)
+        DynamicPatch = DynamicUnionPatch
+        patch._patch = DynamicPatch.deserialize_from_compact_protocol(safe_patch.data)
         return patch
 
 
@@ -397,7 +413,8 @@ def to_safe_patch(self) -> _fbthrift_safe_patch_types.MyUnionSafePatch:
     @staticmethod
     def from_safe_patch(safe_patch: _fbthrift_safe_patch_types.MyUnionSafePatch) -> MyUnionPatch:
         patch = MyUnionPatch()
-        patch._patch = DynamicStructPatch.deserialize_from_compact_protocol(safe_patch.data)
+        DynamicPatch = DynamicUnionPatch
+        patch._patch = DynamicPatch.deserialize_from_compact_protocol(safe_patch.data)
         return patch
 
 
@@ -420,7 +437,8 @@ def to_safe_patch(self) -> _fbthrift_safe_patch_types.MyUnionSafePatch:
     @staticmethod
     def from_safe_patch(safe_patch: _fbthrift_safe_patch_types.MyUnionSafePatch) -> MyUnionPatch:
         patch = MyUnionPatch()
-        patch._patch = DynamicStructPatch.deserialize_from_compact_protocol(safe_patch.data)
+        DynamicPatch = DynamicUnionPatch
+        patch._patch = DynamicPatch.deserialize_from_compact_protocol(safe_patch.data)
         return patch
 
 
@@ -449,7 +467,8 @@ def to_safe_patch(self) -> _fbthrift_safe_patch_types.MyExceptionSafePatch:
     @staticmethod
     def from_safe_patch(safe_patch: _fbthrift_safe_patch_types.MyExceptionSafePatch) -> MyExceptionPatch:
         patch = MyExceptionPatch()
-        patch._patch = DynamicStructPatch.deserialize_from_compact_protocol(safe_patch.data)
+        DynamicPatch = DynamicStructPatch
+        patch._patch = DynamicPatch.deserialize_from_compact_protocol(safe_patch.data)
         return patch
 
 
@@ -472,7 +491,8 @@ def to_safe_patch(self) -> _fbthrift_safe_patch_types.MyExceptionSafePatch:
     @staticmethod
     def from_safe_patch(safe_patch: _fbthrift_safe_patch_types.MyExceptionSafePatch) -> MyExceptionPatch:
         patch = MyExceptionPatch()
-        patch._patch = DynamicStructPatch.deserialize_from_compact_protocol(safe_patch.data)
+        DynamicPatch = DynamicStructPatch
+        patch._patch = DynamicPatch.deserialize_from_compact_protocol(safe_patch.data)
         return patch
 
 
@@ -495,7 +515,8 @@ def to_safe_patch(self) -> _fbthrift_safe_patch_types.MyExceptionSafePatch:
     @staticmethod
     def from_safe_patch(safe_patch: _fbthrift_safe_patch_types.MyExceptionSafePatch) -> MyExceptionPatch:
         patch = MyExceptionPatch()
-        patch._patch = DynamicStructPatch.deserialize_from_compact_protocol(safe_patch.data)
+        DynamicPatch = DynamicStructPatch
+        patch._patch = DynamicPatch.deserialize_from_compact_protocol(safe_patch.data)
         return patch
 
 
@@ -518,7 +539,8 @@ def to_safe_patch(self) -> _fbthrift_safe_patch_types.MyExceptionSafePatch:
     @staticmethod
     def from_safe_patch(safe_patch: _fbthrift_safe_patch_types.MyExceptionSafePatch) -> MyExceptionPatch:
         patch = MyExceptionPatch()
-        patch._patch = DynamicStructPatch.deserialize_from_compact_protocol(safe_patch.data)
+        DynamicPatch = DynamicStructPatch
+        patch._patch = DynamicPatch.deserialize_from_compact_protocol(safe_patch.data)
         return patch
 
 
@@ -547,7 +569,8 @@ def to_safe_patch(self) -> _fbthrift_safe_patch_types.MyExceptionWithMessageSafe
     @staticmethod
     def from_safe_patch(safe_patch: _fbthrift_safe_patch_types.MyExceptionWithMessageSafePatch) -> MyExceptionWithMessagePatch:
         patch = MyExceptionWithMessagePatch()
-        patch._patch = DynamicStructPatch.deserialize_from_compact_protocol(safe_patch.data)
+        DynamicPatch = DynamicStructPatch
+        patch._patch = DynamicPatch.deserialize_from_compact_protocol(safe_patch.data)
         return patch
 
 
@@ -570,7 +593,8 @@ def to_safe_patch(self) -> _fbthrift_safe_patch_types.MyExceptionWithMessageSafe
     @staticmethod
     def from_safe_patch(safe_patch: _fbthrift_safe_patch_types.MyExceptionWithMessageSafePatch) -> MyExceptionWithMessagePatch:
         patch = MyExceptionWithMessagePatch()
-        patch._patch = DynamicStructPatch.deserialize_from_compact_protocol(safe_patch.data)
+        DynamicPatch = DynamicStructPatch
+        patch._patch = DynamicPatch.deserialize_from_compact_protocol(safe_patch.data)
         return patch
 
 
@@ -593,7 +617,8 @@ def to_safe_patch(self) -> _fbthrift_safe_patch_types.MyExceptionWithMessageSafe
     @staticmethod
     def from_safe_patch(safe_patch: _fbthrift_safe_patch_types.MyExceptionWithMessageSafePatch) -> MyExceptionWithMessagePatch:
         patch = MyExceptionWithMessagePatch()
-        patch._patch = DynamicStructPatch.deserialize_from_compact_protocol(safe_patch.data)
+        DynamicPatch = DynamicStructPatch
+        patch._patch = DynamicPatch.deserialize_from_compact_protocol(safe_patch.data)
         return patch
 
 
@@ -616,7 +641,8 @@ def to_safe_patch(self) -> _fbthrift_safe_patch_types.MyExceptionWithMessageSafe
     @staticmethod
     def from_safe_patch(safe_patch: _fbthrift_safe_patch_types.MyExceptionWithMessageSafePatch) -> MyExceptionWithMessagePatch:
         patch = MyExceptionWithMessagePatch()
-        patch._patch = DynamicStructPatch.deserialize_from_compact_protocol(safe_patch.data)
+        DynamicPatch = DynamicStructPatch
+        patch._patch = DynamicPatch.deserialize_from_compact_protocol(safe_patch.data)
         return patch
 
 
@@ -645,13 +671,14 @@ def to_safe_patch(self) -> _fbthrift_safe_patch_types.ReservedKeywordSafePatch:
     @staticmethod
     def from_safe_patch(safe_patch: _fbthrift_safe_patch_types.ReservedKeywordSafePatch) -> ReservedKeywordPatch:
         patch = ReservedKeywordPatch()
-        patch._patch = DynamicStructPatch.deserialize_from_compact_protocol(safe_patch.data)
+        DynamicPatch = DynamicStructPatch
+        patch._patch = DynamicPatch.deserialize_from_compact_protocol(safe_patch.data)
         return patch
 
 
 
 class UnionToBeRenamedPatch(
-    BaseStructPatch[_fbthrift__test__fixtures__basic__module__thrift_types.UnionToBeRenamed]
+    BaseUnionPatch[_fbthrift__test__fixtures__basic__module__thrift_types.UnionToBeRenamed]
 ):
     pass
 
@@ -674,7 +701,8 @@ def to_safe_patch(self) -> _fbthrift_safe_patch_types.UnionToBeRenamedSafePatch:
     @staticmethod
     def from_safe_patch(safe_patch: _fbthrift_safe_patch_types.UnionToBeRenamedSafePatch) -> UnionToBeRenamedPatch:
         patch = UnionToBeRenamedPatch()
-        patch._patch = DynamicStructPatch.deserialize_from_compact_protocol(safe_patch.data)
+        DynamicPatch = DynamicUnionPatch
+        patch._patch = DynamicPatch.deserialize_from_compact_protocol(safe_patch.data)
         return patch