Add MutableStructOrUnion.fbthrift_reset()

Summary: Currently only implemented for `MutableUnion`s. The implementation for `MutableStruct`s is forthcoming.

Reviewed By: createdbysk

Differential Revision: D66901416

fbshipit-source-id: b764bed4bcf30825a86ca285ae4c8d0294858345

third-party/thrift/src/thrift/lib/python/mutable_types.pyi

@@ -20,15 +20,16 @@ from thrift.python.exceptions import Error
 
 # Base class for mutable structs and mutable unions
 class MutableStructOrUnion:
+    def fbthrift_reset(self) -> None: ...
     def __eq__(self, other: object) -> bool: ...
     def __lt__(self, other: object) -> bool: ...
     def __le__(self, other: object) -> bool: ...
 
 class MutableStructMeta(type): ...
-class MutableStruct(MutableStructOrUnion, metaclass=MutableStructMeta): ...
 class MutableUnionMeta(type): ...
 class MutableGeneratedError(Error): ...
-class MutableUnion(MutableStructOrUnion): ...
+class MutableStruct(MutableStructOrUnion, metaclass=MutableStructMeta): ...
+class MutableUnion(MutableStructOrUnion, metaclass=MutableUnionMeta): ...
 
 MutableStructOrError = typing.Union[MutableStruct, MutableGeneratedError]
 

third-party/thrift/src/thrift/lib/python/mutable_types.pyx

@@ -232,10 +232,18 @@ cdef void set_mutable_struct_field(list struct_list, int16_t index, value) excep
 
 
 cdef class MutableStructOrUnion:
+    def fbthrift_reset(self):
+        """
+        Resets this Thrift struct or union instance to its (standard) default values.
+        """
+        raise NotImplementedError("Not implemented on base MutableStructOrUnion class")
+
     cdef IOBuf _fbthrift_serialize(self, Protocol proto):
         raise NotImplementedError("Not implemented on base MutableStructOrUnion class")
+
     cdef uint32_t _fbthrift_deserialize(self, IOBuf buf, Protocol proto) except? 0:
         raise NotImplementedError("Not implemented on base MutableStructOrUnion class")
+
     cdef _fbthrift_get_field_value(self, int16_t index):
         raise NotImplementedError("Not implemented on base MutableStructOrUnion class")
 
@@ -290,6 +298,12 @@ cdef class MutableStruct(MutableStructOrUnion):
     def __init__(self, **kwargs):
         pass
 
+    def fbthrift_reset(self):
+        raise NotImplementedError(
+            "fbthrift_reset() is not (yet) implemented for mutable thrift-python "
+            "Struct types."
+        )
+
     def __call__(self, **kwargs):
         self_copy = copy.deepcopy(self)
 
@@ -843,6 +857,17 @@ cdef class MutableUnion(MutableStructOrUnion):
 
         self._fbthrift_set_mutable_union_value(field_id, field_python_value)
 
+    def fbthrift_reset(self):
+        """
+        Resets (or "clears") this union, ensuring none of its fields are set.
+
+        Following this call, the following are true:
+            * the current value (`self.fbthrift_current_value`) is`None`
+            * the current field (`self.fbthrift_current_field`) is the (special)
+              `EMPTY` value.
+        """
+        self._fbthrift_set_mutable_union_value(field_id=0, field_python_value=None)
+
     cdef void _fbthrift_set_mutable_union_value(
         self, int field_id, object field_python_value
     ) except *:
@@ -943,7 +968,7 @@ cdef class MutableUnion(MutableStructOrUnion):
     ):
         """
         Returns the current value for this union, in "python" format (as opposed to
-        "internal data", see `TypeInfoBase`).
+        "internal data", see `TypeInfoBase`), or None if this union is empty.
 
         This method DOES NOT handle adapted fields, i.e. if the current field has an
         adapter, the returned value will NOT be the adapted value, but rather that of
@@ -952,7 +977,13 @@ cdef class MutableUnion(MutableStructOrUnion):
         Args:
             current_field_enum_value: the field ID of the current field, read from
                 `self._fbthrift_data[0]` and passed in to avoid unnecessarily reading it
-                again.
+                again. This can be 0 (for an empty union).
+
+        Returns:
+            Python-value of the current field, or None if `current_field_enum_value` is
+            0 (i.e., empty union). In the latter case, also asserts that
+            `self._fbthrift_data[1]` (i.e., the current value in internal data
+            representation) is None.
         """
         field_internal_data = self._fbthrift_data[1]
 

third-party/thrift/src/thrift/test/thrift-python/struct_test.py

@@ -492,6 +492,21 @@ def test_default_values(self) -> None:
             TestStructWithDefaultValuesMutable().unqualified_struct_intrinsic_default,
         )
 
+    def test_reset(self) -> None:
+        w = TestStructWithDefaultValuesMutable(optional_integer=123)
+        self.assertEqual(w.optional_integer, 123)
+        self.assertEqual(w.unqualified_integer, 42)
+        with self.assertRaisesRegex(
+            NotImplementedError,
+            (
+                r"fbthrift_reset\(\) is not \(yet\) implemented for mutable "
+                "thrift-python Struct types."
+            ),
+        ):
+            w.fbthrift_reset()
+            # self.assertIsNone(w.optional_integer)
+            # self.assertEqual(w.unqualified_integer, 42)
+
     def test_equality_and_hashability(self) -> None:
         # Equality
         w_mutable = TestStructMutable(unqualified_string="hello, world!")

third-party/thrift/src/thrift/test/thrift-python/union_test.py

@@ -795,6 +795,38 @@ def test_set_nonexistent_field(self) -> None:
             # pyre-ignore[16]: Intentional for test
             u.non_existent_field = 999
 
+    def test_reset(self) -> None:
+        u = TestUnionMutable(string_field="Hello!")
+        self.assertEqual(u.fbthrift_current_value, "Hello!")
+        self.assertIs(
+            u.fbthrift_current_field,
+            TestUnionMutable.FbThriftUnionFieldEnum.string_field,
+        )
+
+        # Resetting union makes it empty
+        u.fbthrift_reset()
+        self.assertIsNone(u.fbthrift_current_value)
+        self.assertIs(
+            u.fbthrift_current_field,
+            TestUnionMutable.FbThriftUnionFieldEnum.EMPTY,
+        )
+
+        # Union can be reset again, it's a no-op
+        u.fbthrift_reset()
+        self.assertIsNone(u.fbthrift_current_value)
+
+        # Union can then be set:
+        u.int_field = 42
+        self.assertEqual(u.fbthrift_current_value, 42)
+        self.assertIs(
+            u.fbthrift_current_field,
+            TestUnionMutable.FbThriftUnionFieldEnum.int_field,
+        )
+
+        # One last reset
+        u.fbthrift_reset()
+        self.assertIsNone(u.fbthrift_current_value)
+
     def test_del_field(self) -> None:
         u = TestUnionMutable(string_field="Hello!")
         self.assertEqual(u.string_field, "Hello!")