Modify function.py to support standard containers.

Modify utils.py to add more tests for when cpp_exact_type is faulty.

C++ functions with return parameters require the code generator to generate a C++ lambda function.

PiperOrigin-RevId: 363927286

Author: kaviram

clif/pybind11/function.py

@@ -38,6 +38,11 @@ def generate_from(module_name: str, func_decl: ast_pb2.FuncDecl,
     pybind11 function bindings code.
   """
 
+  if len(func_decl.returns) >= 2 or (len(func_decl.returns) >= 1 and
+                                     func_decl.cpp_void_return):
+    yield from _generate_return_args_lambda(func_decl, module_name)
+    return
+
   cpp_lambda_return_type = _has_bytes_return(func_decl)
   if cpp_lambda_return_type:
     yield from _generate_cpp_lambda(func_decl, cpp_lambda_return_type,
@@ -86,7 +91,7 @@ def _generate_cpp_function_cast(func_decl: ast_pb2.FuncDecl,
   params_list_types = []
   for param in func_decl.params:
     if param.HasField('cpp_exact_type'):
-      if utils.is_nested_template(param.cpp_exact_type):
+      if not utils.is_usable_cpp_exact_type(param.cpp_exact_type):
         params_list_types.append(param.type.cpp_type)
       else:
         params_list_types.append(param.cpp_exact_type)
@@ -156,7 +161,7 @@ def get_params_strings(func: ast_pb2.FuncDecl):
     lang_types.append(param.type.lang_type)
     cpp_names.append(param.name.cpp_name)
     params_str_with_types_list.append(
-        f'{param.type.lang_type} {param.name.cpp_name}')
+        f'{param.type.cpp_type} {param.name.cpp_name}')
     if param.default_value:
       default_values.append(
           f'py::arg("{param.name.cpp_name}") = {param.default_value}')
@@ -202,3 +207,39 @@ def _generate_cpp_lambda(func_decl: ast_pb2.FuncDecl, return_type: str,
   yield I + I + '}'
   yield I + ');'
   return
+
+
+def _generate_return_args_lambda(func_decl: ast_pb2.FuncDecl, module_name: str):
+  """Generates C++ lambda functions with return parameters."""
+  params_strings = get_params_strings(func_decl)
+  yield (f'{module_name}.def("{func_decl.name.native}",'
+         f'[]({params_strings.names_with_types}) {{')
+
+  main_return = ''
+  main_return_cpp_name = ''
+  other_returns_cpp_names = []
+  for i, r in enumerate(func_decl.returns):
+    if i == 0:
+      main_return_cpp_name = r.name.cpp_name
+      main_return = f'{r.type.cpp_type} {r.name.cpp_name}'
+      continue
+    other_returns_cpp_names.append(r.name.cpp_name)
+    yield I + f'{r.type.cpp_type} {r.name.cpp_name};'
+  other_returns = ', '.join(other_returns_cpp_names)
+  other_returns_params_list = [f'&{r}' for r in other_returns_cpp_names]
+
+  if not func_decl.cpp_void_return:
+    yield I + (f'{main_return} = {func_decl.name.cpp_name}'
+               f'({params_strings.cpp_names}, &y);')
+    yield I + f'return std::make_tuple({main_return_cpp_name}, {other_returns});'
+  else:
+    yield I + f'{main_return};'
+    if not other_returns_cpp_names:
+      yield I + (f'{func_decl.name.cpp_name}({params_strings.cpp_names},'
+                 f'&{main_return_cpp_name});')
+      yield I + f'return {main_return_cpp_name};'
+    else:
+      yield I + (f'{func_decl.name.cpp_name}({params_strings.cpp_names},'
+                 f'&{main_return_cpp_name}, {other_returns_params_list});')
+      yield I + f'return std::make_tuple({main_return_cpp_name}, {other_returns});'
+  yield '});'

clif/pybind11/utils.py

@@ -65,6 +65,10 @@ def is_nested_template(s: str) -> bool:
   return '>>' in s
 
 
+def is_usable_cpp_exact_type(s: str) -> bool:
+  return not is_nested_template(s) or '&' in s
+
+
 # Dataclass used to group together and pass around parameter lists.
 @dataclasses.dataclass
 class ParamsStrings:

clif/testing/python/std_containers_test.py

@@ -14,72 +14,90 @@
 
 """Tests for clif.testing.python.std_containers."""
 
-import unittest
-from clif.testing.python import std_containers
-
-
-class StdContainersTest(unittest.TestCase):
-
-  def testVector(self):
-    self.assertEqual(std_containers.Mul([1, 2, 3], 2), [2, 4, 6])
-
-  def testArray(self):
-    self.assertEqual(std_containers.Div([2, 4], 2), [1, 2])
+from absl.testing import absltest
+from absl.testing import parameterized
 
+from clif.testing.python import std_containers
+# TODO: Restore simple import after OSS setup includes pybind11.
+# pylint: disable=g-import-not-at-top
+try:
+  from clif.testing.python import std_containers_pybind11
+except ImportError:
+  std_containers_pybind11 = None
+# pylint: enable=g-import-not-at-top
+
+
+@parameterized.named_parameters([
+    np for np in zip(('c_api', 'pybind11'), (std_containers,
+                                             std_containers_pybind11))
+    if np[1] is not None
+])
+class StdContainersTest(absltest.TestCase):
+
+  def testVector(self, wrapper_lib):
+    self.assertEqual(wrapper_lib.Mul([1, 2, 3], 2), [2, 4, 6])
+
+  def testArray(self, wrapper_lib):
+    self.assertEqual(wrapper_lib.Div([2, 4], 2), [1, 2])
+
+    error = ValueError
+    if wrapper_lib is std_containers_pybind11:
+      error = TypeError
     # Exceed bounds of array.
-    self.assertRaises(ValueError, lambda: std_containers.Div([2, 4, 6], 2))
+    self.assertRaises(error, lambda: wrapper_lib.Div([2, 4, 6], 2))
 
-  def testVectorBool(self):
-    self.assertEqual(std_containers.Odd([1, 2, 3]), [True, False, True])
-    self.assertEqual(std_containers.Even([1, 2, 3]), [False, True, False])
+  def testVectorBool(self, wrapper_lib):
+    self.assertEqual(wrapper_lib.Odd([1, 2, 3]), [True, False, True])
+    self.assertEqual(wrapper_lib.Even([1, 2, 3]), [False, True, False])
 
-  def testMap(self):
+  def testMap(self, wrapper_lib):
     # pylint: disable=bad-whitespace
-    self.assertEqual(std_containers.Find(0, {1: 2, 3: 4, 5: 0}), (True, 5))
-    r = std_containers.Find(6, {1: 2, 3: 4, 5: 0})
+    self.assertEqual(wrapper_lib.Find(0, {1: 2, 3: 4, 5: 0}), (True, 5))
+    r = wrapper_lib.Find(6, {1: 2, 3: 4, 5: 0})
     self.assertIsInstance(r, tuple)
     self.assertIs(r[0], False)
 
-  def testOnes(self):
-    ones = std_containers.Ones(4, 5)
-    self.assertEqual(len(ones), 4)
+  def testOnes(self, wrapper_lib):
+    ones = wrapper_lib.Ones(4, 5)
+    self.assertLen(ones, 4)
     for row in ones:
-      self.assertEqual(len(row), 5)
+      self.assertLen(row, 5)
       self.assertTrue(all(elem == 1 for elem in row))
 
-  def testCapitals(self):
-    capitals = std_containers.Capitals()
-    self.assertEqual(len(capitals), 3)
+  def testCapitals(self, wrapper_lib):
+    capitals = wrapper_lib.Capitals()
+    self.assertLen(capitals, 3)
     self.assertEqual(capitals[0], ('CA', 'Sacramento'))
     self.assertEqual(capitals[1], ('OR', 'Salem'))
     self.assertEqual(capitals[2], ('WA', 'Olympia'))
 
-  def testMatrixSum(self):
+  def testMatrixSum(self, wrapper_lib):
     a = [[1, 2, 3], [3, 2, 1]]
     b = [[4, 3, 2], [2, 3, 4]]
-    s = std_containers.MatrixSum(a, b)
-    self.assertEqual(len(s), 2)
+    s = wrapper_lib.MatrixSum(a, b)
+    self.assertLen(s, 2)
     for row in s:
-      self.assertEqual(len(row), 3)
+      self.assertLen(row, 3)
       self.assertTrue(all(elem == 5 for elem in row))
 
-  def testMake2By3(self):
-    m = std_containers.Make2By3()
-    self.assertEqual(len(m), 2)
+  def testMake2By3(self, wrapper_lib):
+    m = wrapper_lib.Make2By3()
+    self.assertLen(m, 2)
     for c in m:
-      self.assertEqual(len(c), 3)
+      self.assertLen(c, 3)
 
-  def testFlatten2By3(self):
+  def testFlatten2By3(self, wrapper_lib):
     m = ((1, 2, 3), (4, 5, 6))
-    self.assertEqual(len(m), 2)
-    t = std_containers.Flatten2By3(m)
-    self.assertEqual(len(t), 6)
+    self.assertLen(m, 2)
+    t = wrapper_lib.Flatten2By3(m)
+    self.assertLen(t, 6)
     for e, i in zip(t, range(1, 7)):
       self.assertEqual(e, i)
 
-  def testTakeVectorOfStrings(self):
-    self.assertEqual(std_containers.LastStringInVector(['hello', 'world']),
-                     'world')
+  def testTakeVectorOfStrings(self, wrapper_lib):
+    self.assertEqual(
+        wrapper_lib.LastStringInVector(['hello', 'world']), 'world')
+
 
 if __name__ == '__main__':
-  unittest.main()
+  absltest.main()