Also generate Clif_PyObj functions for protobufs.

Many user defined `Clif_PyObj` functions depend on `Clif_PyObj`functions of protobufs. For example, https://source.corp.google.com/piper///depot/google3/devtools/clif/contrib/python/protobuf/repeated_field.h;l=24.

PiperOrigin-RevId: 494856800

Author: wangxf

clif/pybind11/gen_type_info.py

@@ -216,3 +216,87 @@ class CapsuleType(BaseType):
   def generate_clif_use(self) -> Generator[str, None, None]:
     yield (f'// CLIF use `{self.cpp_name}` as {self.py_name}, '
            'PythonCapsule, Pybind11Ignore')
+
+
+@dataclasses.dataclass
+class ProtoType(BaseType):
+  """Wraps a C++ proto Message."""
+
+  def generate_header(self) -> Generator[str, None, None]:
+    yield ''
+    yield from self.generate_clif_use()
+    yield (f'PyObject* Clif_PyObjFrom(const {self.cpp_name}&, '
+           '::clif::py::PostConv);')
+    yield (f'PyObject* Clif_PyObjFrom(std::shared_ptr<const {self.cpp_name}>,'
+           '::clif::py::PostConv);')
+    yield (f'PyObject* Clif_PyObjFrom(std::unique_ptr<const {self.cpp_name}>,'
+           '::clif::py::PostConv);')
+    yield ''
+    yield f'bool Clif_PyObjAs(PyObject* input, {self.cpp_name}* output);'
+    yield ('bool Clif_PyObjAs(PyObject* input, '
+           f'std::unique_ptr<{self.cpp_name}>* output);')
+
+  def generate_converters(self) -> Generator[str, None, None]:
+    yield ''
+    yield (f'PyObject* Clif_PyObjFrom(const {self.cpp_name}& c, '
+           '::clif::py::PostConv) {')
+    yield I + 'return pybind11::cast(c).release().ptr();'
+    yield '}'
+    yield (f'PyObject* Clif_PyObjFrom(std::shared_ptr<const {self.cpp_name}> c'
+           ', ::clif::py::PostConv) {')
+    yield I + 'return pybind11::cast(std::move(c)).release().ptr();'
+    yield '}'
+    yield ''
+    yield (f'PyObject* Clif_PyObjFrom(std::unique_ptr<const {self.cpp_name}> c'
+           ', ::clif::py::PostConv) {')
+    yield I + 'return pybind11::cast(std::move(c)).release().ptr();'
+    yield '}'
+    yield ''
+    yield f'bool Clif_PyObjAs(PyObject* input, {self.cpp_name}* output) {{'
+    yield I + 'try {'
+    yield I + I + (f'*output = pybind11::cast<{self.cpp_name}>'
+                   '(pybind11::handle(input));')
+    yield I + '} catch (pybind11::cast_error) {'
+    yield I + I + 'return false;'
+    yield I + '}'
+    yield I + 'return true;'
+    yield '}'
+    yield ''
+    yield ('bool Clif_PyObjAs(PyObject* input, '
+           f'std::unique_ptr<{self.cpp_name}>* output) {{')
+    yield I + 'try {'
+    yield I + I + ('*output = pybind11::cast<std::unique_ptr'
+                   f'<{self.cpp_name}>>(pybind11::handle(input));')
+    yield I + '} catch (pybind11::cast_error) {'
+    yield I + I + 'return false;'
+    yield I + '}'
+    yield I + 'return true;'
+    yield '}'
+    yield ''
+
+
+@dataclasses.dataclass
+class ProtoEnumType(BaseType):
+  """Wraps a C++ proto Enum."""
+
+  def generate_header(self) -> Generator[str, None, None]:
+    yield ''
+    yield from self.generate_clif_use()
+    yield f'PyObject* Clif_PyObjFrom({self.cpp_name}, ::clif::py::PostConv);'
+    yield f'bool Clif_PyObjAs(PyObject* input, {self.cpp_name}* output);'
+
+  def generate_converters(self) -> Generator[str, None, None]:
+    yield ''
+    yield (f'PyObject* Clif_PyObjFrom({self.cpp_name} c, '
+           '::clif::py::PostConv) {')
+    yield I + 'return pybind11::cast(c).release().ptr();'
+    yield '}'
+    yield f'bool Clif_PyObjAs(PyObject* input, {self.cpp_name}* output) {{'
+    yield I + 'try {'
+    yield I + I + (f'*output = pybind11::cast<{self.cpp_name}>'
+                   '(pybind11::handle(input));')
+    yield I + '} catch (pybind11::cast_error) {'
+    yield I + I + 'return false;'
+    yield I + '}'
+    yield I + 'return true;'
+    yield '}'

clif/pybind11/generator.py

@@ -248,14 +248,16 @@ def _generate_headlines(self):
     yield '// potential future optimization: generate this line only as needed.'
     yield '#include "third_party/pybind11/include/pybind11/stl.h"'
     yield ''
+    yield '// See pybind11_protobuf/proto_caster_impl.h'
+    yield '#if !defined(PYBIND11_PROTOBUF_UNSAFE)'
+    yield I + '#define PYBIND11_PROTOBUF_UNSAFE 1'
+    yield '#endif'
+    yield ''
     for include in includes:
       yield f'#include "{include}"'
     yield f'#include "{self._header_path}"'
     yield '#include "clif/pybind11/runtime.h"'
     yield '#include "clif/pybind11/type_casters.h"'
-
-    yield '// See pybind11_protobuf/proto_caster_impl.h'
-    yield '#define PYBIND11_PROTOBUF_UNSAFE 1'
     yield '#include "third_party/pybind11_protobuf/native_proto_caster.h"'
     yield ''
     yield 'namespace py = pybind11;'

clif/python/proto.py

@@ -24,6 +24,7 @@
 import argparse
 import itertools
 import sys
+from clif.pybind11 import gen_type_info
 from clif.python import gen
 from clif.python import clif_types as types
 from clif.python.utils import proto_util
@@ -94,17 +95,56 @@ def CreatePyTypeInfo(desc, path,
   return messages
 
 
-def GenerateForPybind11(messages, proto_hdr):
+def GenerateForPybind11(messages, clif_hdr, proto_hdr):
   """Generate no-op header files to bypass the checks of PyCLIF."""
+  proto_types = []
+  for m in messages:
+    ctor = None
+    if isinstance(m, types.ProtoType):
+      ctor = gen_type_info.ProtoType
+    elif isinstance(m, types.CapsuleType):
+      ctor = gen_type_info.CapsuleType
+    elif isinstance(m, types.ProtoEnumType):
+      ctor = gen_type_info.ProtoEnumType
+    if not ctor:
+      raise _ParseError('Unsupported proto type.')
+    proto_type = ctor(
+        cpp_name=m.cname, py_name=m.pyname,
+        cpp_namespace=m.namespace)
+    proto_types.append(proto_type)
+  proto_types = sorted(proto_types, key=lambda gen_type: gen_type.cpp_name)
   with open(FLAGS.header_out, 'w') as hout:
-    hout.write(f'#include "{proto_hdr}"')
-    hout.write('\n')
-    for m in messages:
-      hout.write(f'// CLIF use `{m.cname}` as {m.pyname}, Pybind11Ignore')
-      hout.write('\n')
+    gen.WriteTo(
+        hout,
+        [f'#include "{proto_hdr}"',
+         '#include "clif/python/postconv.h"',
+         '#include "third_party/pybind11/include/pybind11/smart_holder.h"',
+         '#include "third_party/pybind11_protobuf/native_proto_caster.h"'])
+    for namespace, typedefs in itertools.groupby(
+        proto_types, lambda gen_type: gen_type.cpp_namespace):
+      namespace = namespace.strip(':') or 'clif'
+      gen.WriteTo(
+          hout,
+          [' '.join('namespace %s {' % ns for ns in namespace.split('::'))])
+      for t in typedefs:
+        gen.WriteTo(hout, t.generate_header())
+      gen.WriteTo(
+          hout,
+          ['} ' * (1 + namespace.count('::')) + ' // namespace ' + namespace])
 
   with open(FLAGS.ccdeps_out, 'w') as cout:
-    cout.write('')
+    gen.WriteTo(cout, [f'#include "{clif_hdr}"'])
+    for namespace, typedefs in itertools.groupby(
+        proto_types, lambda gen_type: gen_type.cpp_namespace):
+      namespace = namespace.strip(':') or 'clif'
+      gen.WriteTo(
+          cout,
+          [' '.join('namespace %s {' % ns for ns in namespace.split('::'))])
+      for t in typedefs:
+        gen.WriteTo(cout, t.generate_converters())
+      gen.WriteTo(
+          cout,
+          ['} ' * (1 + namespace.count('::')) + ' // namespace ' + namespace])
 
 
 def GenerateFrom(messages, proto_filename, clif_hdr, proto_hdr):
@@ -167,7 +207,7 @@ def main(_):
   if FLAGS.pyclif_codegen_mode == 'c_api':
     GenerateFrom(messages, name, hdr, pypath+'.pb.h')
   else:
-    GenerateForPybind11(messages, pypath+'.pb.h')
+    GenerateForPybind11(messages, hdr, pypath+'.pb.h')
 
 
 def ParseFlags():