From 6203d65e4c75e515dd61b140020a674cafd7ed3e Mon Sep 17 00:00:00 2001
From: Andrew Hilger <ahilger@meta.com>
Date: Tue, 17 Dec 2024 12:23:31 -0800
Subject: [PATCH] make Union.type enum creation lazy

Summary:
Makes `union.Type` enum creation lazy by deferring creation until first access.

NOTE: This requires an interface change where `u.type` now *always* returns the enum even if there's a field name conflict. To safely land, we probably have to add a compile-time validator against this.

Reviewed By: Filip-F, createdbysk

Differential Revision: D67166249

fbshipit-source-id: bfc14f0e95762c93f1fc27141d07ac02ac2624be
---
 .../thrift/src/thrift/lib/python/types.pxd    |  3 +-
 .../thrift/src/thrift/lib/python/types.pyx    | 29 +++++++++++++++----
 .../thrift/test/thrift-python/union_test.py   |  4 +--
 3 files changed, 26 insertions(+), 10 deletions(-)

diff --git a/third-party/thrift/src/thrift/lib/python/types.pxd b/third-party/thrift/src/thrift/lib/python/types.pxd
index 9b964c2e2eb93..477d3c1e6906e 100644
--- a/third-party/thrift/src/thrift/lib/python/types.pxd
+++ b/third-party/thrift/src/thrift/lib/python/types.pxd
@@ -254,7 +254,7 @@ cdef tuple _validate_union_init_kwargs(
 )
 
 cdef class Union(StructOrUnion):
-    cdef readonly object type
+    cdef object py_type
     cdef readonly object value
     cdef void _fbthrift_update_current_field_attributes(self) except *
     cdef object _fbthrift_to_internal_data(self, type_value, value)
@@ -263,6 +263,7 @@ cdef class Union(StructOrUnion):
     cdef uint32_t _deserialize(
         Union self, folly.iobuf.IOBuf buf, Protocol proto
     ) except? 0
+    cdef object _fbthrift_py_type_enum(self)
 
 cdef class BadEnum:
     cdef object _enum
diff --git a/third-party/thrift/src/thrift/lib/python/types.pyx b/third-party/thrift/src/thrift/lib/python/types.pyx
index 02ff3044324b3..3431eb8609871 100644
--- a/third-party/thrift/src/thrift/lib/python/types.pyx
+++ b/third-party/thrift/src/thrift/lib/python/types.pyx
@@ -1398,6 +1398,7 @@ cdef class Union(StructOrUnion):
     """
     def __cinit__(self):
         self._fbthrift_data = createUnionTuple()
+        self.py_type = None
 
     def __init__(self, **kwargs):
         self_type = type(self)
@@ -1474,10 +1475,11 @@ cdef class Union(StructOrUnion):
 
     cdef void _fbthrift_update_current_field_attributes(self) except *:
         """
-        Updates the `type` and `value` attributes from the internal data tuple
+        Updates the `value` attribute from the internal data tuple
         of this union (`self._fbthrift_data`).
+        Resets `py_type` to None
         """
-        self.type = type(self).Type(self._fbthrift_data[0])
+        self.py_type = None
         val = self._fbthrift_data[1]
         if val is None:
             self.value = None
@@ -1507,16 +1509,31 @@ cdef class Union(StructOrUnion):
         if _fbthrift_get_Union_type_int(self) != field_id:
             # TODO in python 3.10 update this to use name and obj fields
             raise AttributeError(
-                f'Union contains a value of type {self.type.name}, not '
+                f'Union contains a value of type {self.get_type().name}, not '
                 f'{type(self).Type(field_id).name}')
         return self.value
 
+
+    cdef object _fbthrift_py_type_enum(self):
+        '''
+        Initializes self.py_type enum if None.
+        '''
+        if self.py_type is None:
+            self.py_type = type(self).Type(
+                _fbthrift_get_Union_type_int(self)
+            )
+        return self.py_type
+
+    @property
+    def type(Union self not None):
+        return self._fbthrift_py_type_enum()
+
     def get_type(Union self not None):
-        return self.type
+        return self._fbthrift_py_type_enum()
 
     @property
     def fbthrift_current_field(Union self not None):
-        return self.type
+        return self._fbthrift_py_type_enum()
 
     @property
     def fbthrift_current_value(Union self not None):
@@ -1596,7 +1613,7 @@ cdef class Union(StructOrUnion):
         return (_unpickle_union, (type(self), b''.join(self._serialize(Protocol.COMPACT))))
 
 
-cdef inline _fbthrift_get_Union_type_int(Union u):
+cdef inline int _fbthrift_get_Union_type_int(Union u):
     return u._fbthrift_data[0]
 
 cdef _make_fget_struct(i):
diff --git a/third-party/thrift/src/thrift/test/thrift-python/union_test.py b/third-party/thrift/src/thrift/test/thrift-python/union_test.py
index 3c78e0282fd1d..251bfd70afe7c 100644
--- a/third-party/thrift/src/thrift/test/thrift-python/union_test.py
+++ b/third-party/thrift/src/thrift/test/thrift-python/union_test.py
@@ -232,10 +232,8 @@ def test_from_value(self) -> None:
         )
 
     def test_from_value_none(self) -> None:
-        # BAD: fromValue(None) results in a weird state, where u.type is None
-        # rather than EMPTY.
         union_none = TestUnionImmutable.fromValue(None)
-        self.assertIsNone(union_none.type)
+        self.assertEqual(union_none.type, union_none.Type.EMPTY)
         self.assertIsNone(union_none.value)
 
     def test_from_value_ambiguous_int_bool(self) -> None: