Fix type mapping for return types of python-backed extension functions.

 - support setting an expected return type
 - handle mapping dyn as an argument kind

PiperOrigin-RevId: 858753278

Author: jnthntatum

BUILD

@@ -49,7 +49,6 @@ pybind_extension(
         ":status_macros",
         "@com_google_absl//absl/base",
         "@com_google_absl//absl/base:no_destructor",
-        "@com_google_absl//absl/base:nullability",
         "@com_google_absl//absl/container:flat_hash_map",
         "@com_google_absl//absl/functional:function_ref",
         "@com_google_absl//absl/log:absl_check",

custom_ext/custom_ext_test.py

@@ -57,6 +57,7 @@ def _create_activation(self, impl) -> cel.Activation:
                 [cel.Type.STRING],
                 False,
                 impl,
+                return_type=cel.Type.STRING,
             )
         ],
     )
@@ -93,6 +94,7 @@ def test_error_no_matching_overload(self, ext):
                 [cel.Type.STRING, cel.Type.INT],
                 False,
                 lambda _: "¡Hola Mundo!",
+                return_type=cel.Type.STRING,
             )
         ],
     )
@@ -189,6 +191,42 @@ def _lost_in_translation_return_none(arg1: str) -> str:  # pylint: disable=unuse
 def _lost_in_translation_raising_error(text: str) -> str:  # pylint: disable=unused-argument
   raise LookupError("Lost in translation")
 
+TEST_EXPRESSIONS = [
+    ("getOrDefaultReceiver", "{'a': 1, 'b': 2}.getOrDefault('c', 3) == 3"),
+    ("getOrDefault", "getOrDefault({'a': 'z', 'b': 'y'}, 'a', 'x') == 'z'"),
+    ("lerp_int", "lerp(1, 2, 0.5) == 1.5"),
+    ("lerp_uint", "lerp(1u, 2u, 0.5) == 1.5"),
+]
+
+
+class PythonTypeMappingsTest(parameterized.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    self.descriptor_pool = descriptor_pool.Default()
+    self.env = cel.NewEnv(
+        self.descriptor_pool,
+        variables={},
+        extensions=[sample_cel_ext.SampleCelExtension()],
+    )
+
+  def _compile_expr(self, expression: str) -> cel.Expression:
+    return self.env.compile(expression)
+
+  @parameterized.named_parameters(TEST_EXPRESSIONS)
+  def test_expression(self, expr):
+    compiled_expr = self.env.compile(expr)
+    act = self.env.Activation()
+    res = compiled_expr.eval(act)
+    self.assertEqual(res.value(), True)
+
+  def test_lerp_error_out_of_bounds(self):
+    compiled_expr = self.env.compile("lerp(1, 2, 1.5)")
+    act = self.env.Activation()
+    res = compiled_expr.eval(act)
+    self.assertEqual(res.type(), cel.Type.ERROR)
+    self.assertIn("t must be between 0.0 and 1.0", res.value())
+
 
 if __name__ == "__main__":
   absltest.main()

custom_ext/sample_cel_ext.py

@@ -13,6 +13,9 @@
 # limitations under the License.
 
 """A sample CEL extension implemented entirely in Python."""
+
+from typing import Any
+
 import py_cel as cel
 
 
@@ -51,6 +54,61 @@ def __init__(self):
                     )
                 ],
             ),
+            # Tests for type adaptation.
+            cel.FunctionDecl(
+                "lerp",
+                [
+                    cel.Overload(
+                        "lerp_int_int_double",
+                        return_type=cel.Type.DOUBLE,
+                        parameters=[
+                            cel.Type.INT,
+                            cel.Type.INT,
+                            cel.Type.DOUBLE,
+                        ],
+                        is_member=False,
+                        impl=self.lerp,
+                    ),
+                    cel.Overload(
+                        "lerp_uint_uint_double",
+                        return_type=cel.Type.DOUBLE,
+                        parameters=[
+                            cel.Type.UINT,
+                            cel.Type.UINT,
+                            cel.Type.DOUBLE,
+                        ],
+                        is_member=False,
+                        impl=self.lerp,
+                    ),
+                ],
+            ),
+            cel.FunctionDecl(
+                "getOrDefault",
+                [
+                    cel.Overload(
+                        "map_get_or_default_string_dyn",
+                        return_type=cel.Type.DYN,
+                        parameters=[
+                            cel.Type.Map(cel.Type.STRING, cel.Type.DYN),
+                            cel.Type.STRING,
+                            cel.Type.DYN,
+                        ],
+                        is_member=True,
+                        impl=self.map_get_or_default,
+                    ),
+                    cel.Overload(
+                        "get_or_default_map_string_dyn",
+                        return_type=cel.Type.DYN,
+                        parameters=[
+                            cel.Type.Map(cel.Type.STRING, cel.Type.DYN),
+                            cel.Type.STRING,
+                            cel.Type.DYN,
+                        ],
+                        is_member=False,
+                        impl=self.map_get_or_default,
+                    ),
+                ],
+            ),
         ],
     )
 
@@ -60,3 +118,15 @@ def translate(self, text: str, from_lang: str, to_lang: str) -> str:
     if text != "Hello, world!":
       raise ValueError("Come on, this is just 'Hello, world!'")
     return "¡Hola Mundo!"
+
+  def lerp(self, a: int, b: int, t: float) -> float:
+    """Linearly interpolate between a and b using t."""
+    if t < 0.0 or t > 1.0:
+      raise ValueError("t must be between 0.0 and 1.0")
+    return a + (b - a) * t
+
+  def map_get_or_default(
+      self, m: dict[str, Any], key: str, default: Any
+  ) -> Any:
+    """Get the value for the key from the map, or return the default value."""
+    return m.get(key, default)

py_cel_activation.cc

@@ -74,7 +74,8 @@ PyCelActivation::PyCelActivation(
                                             kFunctionDescriptorOptions);
     activation_.InsertFunction(
         func_descriptor, std::make_unique<PyCelFunctionAdapter>(
-                             function->function_name(), function->impl()));
+                             function->function_name(), function->return_type(),
+                             function->impl()));
   }
 };
 

py_cel_function.cc

@@ -56,15 +56,18 @@ static std::shared_ptr<PyCelEnvInternal> GetEnvFromContext(
 
 void PyCelFunction::DefinePythonBindings(pybind11::module& m) {
   py::class_<PyCelFunction, std::shared_ptr<PyCelFunction>>(m, "Function")
-      .def(py::init<std::string, std::vector<PyCelType>, bool, PyObject*>(),
+      .def(py::init<std::string, std::vector<PyCelType>, bool, PyObject*,
+                    PyCelType>(),
            py::arg("function_name"), py::arg("parameters"),
-           py::arg("is_member"), py::arg("impl"));
+           py::arg("is_member"), py::arg("impl"),
+           py::arg("return_type") = PyCelType::Dyn());
 }
 
 PyCelFunction::PyCelFunction(std::string function_name,
                              std::vector<PyCelType> parameters, bool is_member,
-                             PyObject* impl)
+                             PyObject* impl, PyCelType return_type)
     : function_name_(std::move(function_name)),
+      return_type_(std::move(return_type)),
       parameters_(std::move(parameters)),
       is_member_(is_member),
       impl_(impl) {
@@ -79,8 +82,11 @@ PyCelFunction::~PyCelFunction() {
 };
 
 PyCelFunctionAdapter::PyCelFunctionAdapter(std::string function_name,
+                                           PyCelType return_type,
                                            PyObject* py_function)
-    : function_name_(std::move(function_name)), py_function_(py_function) {
+    : function_name_(std::move(function_name)),
+      return_type_(std::move(return_type)),
+      py_function_(py_function) {
   Py_XINCREF(py_function_);
 }
 
@@ -111,7 +117,7 @@ absl::StatusOr<cel::Value> PyCelFunctionAdapter::Invoke(
   }
 
   return PyObjectToCelValue(
-      result, PyCelType::String(),
+      result, return_type_,
       [this]() {
         return absl::StrFormat("Python function '%s'",
                                PyUnicode_AsUTF8(PyObject_Repr(py_function_)));

py_cel_function.h

@@ -17,19 +17,14 @@
 
 #include <Python.h>  // IWYU pragma: keep - Needed for PyObject
 
-#include <memory>
 #include <string>
 #include <vector>
 
-#include "absl/base/nullability.h"
 #include "absl/status/statusor.h"
 #include "absl/types/span.h"
 #include "common/value.h"
 #include "runtime/function.h"
 #include "py_cel_type.h"
-#include "google/protobuf/arena.h"
-#include "google/protobuf/descriptor.h"
-#include "google/protobuf/message.h"
 #include <pybind11/pybind11.h>
 
 namespace cel_python {
@@ -43,16 +38,18 @@ class PyCelFunction {
   static void DefinePythonBindings(pybind11::module& m);
 
   PyCelFunction(std::string function_name, std::vector<PyCelType> parameters,
-                bool is_member, PyObject* impl);
+                bool is_member, PyObject* impl, PyCelType return_type);
   ~PyCelFunction();
 
   std::string function_name() const { return function_name_; }
   const std::vector<PyCelType>& parameters() const { return parameters_; }
   bool is_member() const { return is_member_; }
   PyObject* impl() const { return impl_; }
+  const PyCelType& return_type() const { return return_type_; }
 
  private:
   std::string function_name_;
+  PyCelType return_type_;
   std::vector<PyCelType> parameters_;
   bool is_member_;
   PyObject* impl_;
@@ -62,7 +59,8 @@ class PyCelFunction {
 // function.
 class PyCelFunctionAdapter : public cel::Function {
  public:
-  PyCelFunctionAdapter(std::string function_name, PyObject* py_function);
+  PyCelFunctionAdapter(std::string function_name, PyCelType return_type,
+                       PyObject* py_function);
   ~PyCelFunctionAdapter() override;
 
   absl::StatusOr<cel::Value> Invoke(
@@ -71,6 +69,7 @@ class PyCelFunctionAdapter : public cel::Function {
 
  private:
   std::string function_name_;
+  PyCelType return_type_;
   PyObject* py_function_;
 };
 

py_cel_python_extension.cc

@@ -30,7 +30,6 @@
 #include "compiler/compiler.h"
 #include "runtime/runtime_builder.h"
 #include "runtime/runtime_options.h"
-#include "py_cel_env_internal.h"
 #include "py_cel_extension.h"
 #include "py_cel_function.h"
 #include "py_cel_function_decl.h"
@@ -102,15 +101,22 @@ absl::Status PyCelPythonExtension::ConfigureRuntime(
       std::vector<cel::Kind> types;
       types.reserve(overload.parameters().size());
       for (const PyCelType& arg : overload.parameters()) {
-        types.push_back(arg.GetKind());
+        if (arg.GetKind() == cel::Kind::kDyn) {
+          // C++ runtime dispatcher historically uses kAny for wildcard type (
+          // not distinguishing between dyn and any)
+          types.push_back(cel::Kind::kAny);
+        } else {
+          types.push_back(arg.GetKind());
+        }
       }
 
       cel::FunctionDescriptor descriptor(function.name(), overload.is_member(),
                                          types, kFunctionDescriptorOptions);
       if (overload.py_function()) {
         PY_CEL_RETURN_IF_ERROR(runtime_builder.function_registry().Register(
             descriptor, std::make_unique<PyCelFunctionAdapter>(
-                            function.name(), overload.py_function())));
+                            function.name(), overload.return_type(),
+                            overload.py_function())));
       } else {
         PY_CEL_RETURN_IF_ERROR(
             runtime_builder.function_registry().RegisterLazyFunction(

py_cel_type.h

@@ -41,6 +41,10 @@ class PyCelType {
   // Creates a dynamic type.
   PyCelType();
   PyCelType(const PyCelType& other) = default;
+  PyCelType& operator=(const PyCelType& other) = default;
+  PyCelType(PyCelType&& other) = default;
+  PyCelType& operator=(PyCelType&& other) = default;
+
   // Creates a message type.
   explicit PyCelType(const std::string& name);
   PyCelType(cel::Kind kind, const std::string& name);