fix union IOBufPtr compilation

Summary: The `union_field_ref` is frustrating to work with because on access it always dereferences pointer types. Similarly, it expects the dereferenced pointer element type. This is a problem only  for `std::unique_ptr<folly::IOBuf>` because we define different `Constructor` and `Extractor` specializations for this type only. This is because the `folly.IOBuf` python type holds a `std::unique_ptr`, so it has different methods for extracting the `unique_ptr` vs the `folly::IOBuf` itself.

Reviewed By: Filip-F

Differential Revision: D69155104

fbshipit-source-id: 71987e1c9390b833cb6e07f3c82fab44d56b6a78

thrift/compiler/generate/t_mstch_python_capi_generator.cc

@@ -629,13 +629,19 @@ class python_capi_mstch_field : public mstch_field {
         {
             {"field:cpp_name", &python_capi_mstch_field::cpp_name},
             {"field:marshal_type", &python_capi_mstch_field::marshal_type},
+            {"field:iobuf?", &python_capi_mstch_field::iobuf},
         });
   }
 
   mstch::node cpp_name() { return cpp2::get_name(field_); }
 
   mstch::node marshal_type() { return format_marshal_type(*field_); }
 
+  mstch::node iobuf() {
+    const auto* ttype = field_->get_type()->get_true_type();
+    return ttype->is_binary() && is_type_iobuf(ttype);
+  }
+
  private:
   const std::string py_name_;
 };

thrift/compiler/generate/templates/python_capi/common/union_marshal_type.mustache

@@ -0,0 +1,23 @@
+{{!
+
+  Copyright (c) Meta Platforms, Inc. and affiliates.
+
+  Licensed under the Apache License, Version 2.0 (the "License");
+  you may not use this file except in compliance with the License.
+  You may obtain a copy of the License at
+
+      http://www.apache.org/licenses/LICENSE-2.0
+
+  Unless required by applicable law or agreed to in writing, software
+  distributed under the License is distributed on an "AS IS" BASIS,
+  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  See the License for the specific language governing permissions and
+  limitations under the License.
+
+}}{{!
+
+  The union_field_ref interface makes `std::unique_ptr<folly::IOBuf>` behave like
+  folly::IOBuf, so use the folly::IOBuf Extractor and Constructor instead
+
+}}{{#field:iobuf?}}folly::IOBuf{{/field:iobuf?}}{{!
+}}{{^field:iobuf?}}{{field:marshal_type}}{{/field:iobuf?}}

thrift/compiler/generate/templates/python_capi/thrift_types_capi.cpp.mustache

@@ -121,7 +121,7 @@ Extractor<::apache::thrift::python::capi::ComposedStruct<
       break; // union is unset
   {{#struct:fields}}
     case {{field:key}}:
-      Extractor<{{field:marshal_type}}>{}.extractInto(
+      Extractor<{{> common/union_marshal_type}}>{}.extractInto(
           cpp.{{field:cpp_name}}_ref(), PyTuple_GET_ITEM(fbThriftData, 1), error);
       break;
   {{/struct:fields}}
@@ -249,7 +249,7 @@ PyObject* Constructor<::apache::thrift::python::capi::ComposedStruct<
   {{#struct:fields}}
     case {{field:key}}:
       py_val = StrongRef(
-          Constructor<{{field:marshal_type}}>{}
+          Constructor<{{> common/union_marshal_type}}>{}
           .constructFrom(val.{{field:cpp_name}}_ref()));
       break;
   {{/struct:fields}}

thrift/compiler/test/fixtures/python_capi/out/python_capi/gen-python-capi/module/thrift_types_capi.cpp

@@ -2045,7 +2045,7 @@ Extractor<::apache::thrift::python::capi::ComposedStruct<
           cpp.iobuf_ref(), PyTuple_GET_ITEM(fbThriftData, 1), error);
       break;
     case 2:
-      Extractor<std::unique_ptr<folly::IOBuf>>{}.extractInto(
+      Extractor<folly::IOBuf>{}.extractInto(
           cpp.iobuf_ptr_ref(), PyTuple_GET_ITEM(fbThriftData, 1), error);
       break;
     case 3:
@@ -2107,7 +2107,7 @@ PyObject* Constructor<::apache::thrift::python::capi::ComposedStruct<
       break;
     case 2:
       py_val = StrongRef(
-          Constructor<std::unique_ptr<folly::IOBuf>>{}
+          Constructor<folly::IOBuf>{}
           .constructFrom(val.iobuf_ptr_ref()));
       break;
     case 3:

thrift/test/python_capi/capi_test.py

@@ -44,6 +44,7 @@
     Onion as MyUnion,
     PrimitiveStruct,
     SetStruct,
+    SomeBinary,
     StringPair,
 )
 from thrift.test.python_capi.serialized_dep.thrift_types import (
@@ -229,6 +230,23 @@ def test_roundtrip_union(self) -> None:
         for u in self.my_union():
             self.assertEqual(u, fixture.roundtrip_MyUnion(u))
 
+    def test_roundtrip_iobuf_union(self) -> None:
+        iobuf: IOBuf = IOBuf(memoryview(b"hello world"))
+        checked = 0
+        for field_name in dir(SomeBinary()):
+            if "iobuf" not in field_name:
+                continue
+
+            def make_union(field_name: str) -> SomeBinary:
+                return SomeBinary(**{field_name: iobuf})
+
+            result = fixture.roundtrip_SomeBinary(make_union(field_name))
+            self.assertIsNot(getattr(result, field_name), iobuf)
+            self.assertEqual(getattr(result, field_name), iobuf)
+            self.assertEqual(make_union(field_name), result)
+            checked += 1
+        self.assertEqual(checked, 3)
+
     def test_roundtrip_enum(self) -> None:
         self.assertEqual(MyEnum.MyValue1, fixture.roundtrip_MyEnum(MyEnum.MyValue1))
         self.assertEqual(MyEnum.MyValue2, fixture.roundtrip_MyEnum(MyEnum.MyValue2))

thrift/test/python_capi/fixture.pyi

@@ -35,6 +35,7 @@ from thrift.test.python_capi.module.thrift_types import (
     Onion as MyUnion,
     PrimitiveStruct,
     SetStruct,
+    SomeBinary,
     StringPair,
 )
 from thrift.test.python_capi.serialized_dep.thrift_types import SerializedStruct
@@ -53,6 +54,7 @@ def roundtrip_MapStruct(x: MapStruct) -> MapStruct: ...
 def roundtrip_ComposeStruct(x: ComposeStruct) -> ComposeStruct: ...
 def roundtrip_AdaptedFields(x: AdaptedFields) -> AdaptedFields: ...
 def roundtrip_SerializedStruct(x: SerializedStruct) -> SerializedStruct: ...
+def roundtrip_SomeBinary(x: SomeBinary) -> SomeBinary: ...
 def gen_SerializedStruct(len: int) -> SerializedStruct: ...
 def check_MyStruct(x: object) -> bool: ...
 def check_MyDataItem(x: object) -> bool: ...

thrift/test/python_capi/fixture.pyx

@@ -31,6 +31,7 @@ cdef extern from "thrift/test/python_capi/gen-cpp2/module_types.h" namespace "th
     cppclass SetStruct
     cppclass StringPair
     cppclass VapidStruct
+    cppclass SomeBinary
 
 cdef extern from "thrift/test/python_capi/gen-cpp2/containers_types.h" namespace "thrift::test::python_capi":
     cppclass TemplateLists
@@ -97,6 +98,9 @@ def roundtrip_AdaptedFields(object x):
 def roundtrip_SerializedStruct(object x):
     return __shim__roundtrip[SerializedStruct](x)
 
+def roundtrip_SomeBinary(object x):
+    return __shim__roundtrip[SomeBinary](x)
+
 def gen_SerializedStruct(int64_t len_):
     return __shim__gen_SerializedStruct(len_)
 

thrift/test/python_capi/module.thrift

@@ -208,3 +208,11 @@ union Onion {
   9: map<binary, string> strMap;
   10: id.ProtocolId adapted_int;
 }
+
+union SomeBinary {
+  1: IOBuf iobuf;
+  2: IOBufPtr iobuf_ptr;
+  @cpp.Ref{type = cpp.RefType.Unique}
+  @cpp.AllowLegacyNonOptionalRef
+  3: IOBuf iobufRef;
+}