Add explicit declaration of functions as NewEnv(functions=[...])

PiperOrigin-RevId: 896175602

Author: dmitriplotnikov

README.md

@@ -188,6 +188,76 @@ Resulting message type: <class '...TestAllTypes'>
 Resulting message value: 123
 ```
 
+### Custom functions
+
+When configuring `cel.Env` you can supply custom functions. For each function
+there needs to be a declaration and an implementation. The implementation, which
+is a Python function, can be provided either as part of the declaration
+or separately.
+
+Let's say we want to be able to invoke this function from CEL expressions:
+```python
+  def good_time_of_day(ampm, arg):
+    if ampm == 'am':
+      time_of_day = 'morning'
+    else:
+      time_of_day = 'afternoon'
+    return f"Good {time_of_day}, {arg}"
+```
+
+The implementation can be supplied along with the declaration:
+```python
+  cel_env = cel.NewEnv(functions=[
+            cel.FunctionDecl(
+                "hello",
+                [
+                    cel.Overload(
+                        "hello(string,string)",
+                        return_type=cel.Type.STRING,
+                        parameters=[
+                            cel.Type.STRING,
+                            cel.Type.STRING,
+                        ],
+                        impl=good_time_of_day,
+                    )
+                ],
+            )
+        ])
+```
+
+It can also be provided separately in a dictionary that maps overload IDs to
+their respective implementations:
+
+```python
+  cel_env = cel.NewEnv(functions=[
+            cel.FunctionDecl(
+                "hello",
+                [
+                    cel.Overload(
+                        "hello(string,string)",
+                        return_type=cel.Type.STRING,
+                        parameters=[
+                            cel.Type.STRING,
+                            cel.Type.STRING,
+                        ],
+                    )
+                ],
+            )
+        ],
+        function_impls={
+            "hello(string,string)": good_time_of_day,
+        })
+```
+
+Now that the function implementation is bound to the CEL environment,
+we can invoke it from CEL like this:
+```python
+    result = env.compile("hello('am', 'breakfast is ready!')").eval()
+    print(result.value())   # Good morning, breakfast is ready!
+    result = env.compile("hello('pm', 'tea is served.')").eval()
+    print(result.value())   # Good afternoon, tea is served.
+```
+
 ### Extensions
 
 #### Standard extensions

cel_expr_python/cel_env_test.py

@@ -367,6 +367,105 @@ def test_config_extension_override_different_version(self):
         str(e.exception),
     )
 
+  def test_config_functions(self):
+    config = cel.NewEnvConfigFromYaml("""
+      functions:
+        - name: is_ok
+          overloads:
+            - id: "is_ok_string"
+              target:
+                type_name: string
+              return:
+                type_name: bool
+      """)
+    env = cel.NewEnv(
+        config=config,
+        functions=[
+            cel.FunctionDecl(
+                "hello",
+                [
+                    cel.Overload(
+                        "good_time_of_day",
+                        return_type=cel.Type.STRING,
+                        parameters=[
+                            cel.Type.STRING,
+                            cel.Type.STRING,
+                        ],
+                        impl=lambda ampm, arg: (
+                            "Good"
+                            f" {'morning' if ampm == 'am' else 'afternoon'},"
+                            f" {arg}!"
+                        ),
+                    )
+                ],
+            )
+        ],
+        function_impls={
+            "is_ok_string": lambda arg: arg in ["excellent", "good", "fair"],
+        },
+    )
+    yaml = env.config().to_yaml()
+    self.assertEqual(
+        normalize_yaml(yaml),
+        normalize_yaml("""
+          functions:
+            - name: "hello"
+              overloads:
+                - id: "good_time_of_day"
+                  args:
+                    - type_name: "string"
+                    - type_name: "string"
+                  return:
+                    type_name: "string"
+            - name: "is_ok"
+              overloads:
+                - id: "is_ok_string"
+                  target:
+                    type_name: "string"
+                  return:
+                    type_name: "bool"
+        """),
+    )
+    res = env.compile("hello('am', 'Sunshine')").eval()
+    self.assertEqual(res.value(), "Good morning, Sunshine!")
+    res = env.compile("hello('pm', 'tea is served')").eval()
+    self.assertEqual(res.value(), "Good afternoon, tea is served!")
+    res = env.compile("'good'.is_ok()").eval()
+    self.assertTrue(res.value())
+    res = env.compile("'bad'.is_ok()").eval()
+    self.assertFalse(res.value())
+
+  def test_config_function_override(self):
+    config = cel.NewEnvConfigFromYaml("""
+      functions:
+        - name: foo
+          overloads:
+            - id: "unique_id"
+      """)
+    with self.assertRaises(Exception) as e:
+      cel.NewEnv(
+          config=config,
+          functions=[
+              cel.FunctionDecl(
+                  "bar",
+                  [
+                      cel.Overload(
+                          "unique_id",
+                          impl=lambda: "hello",
+                      )
+                  ],
+              )
+          ],
+          function_impls={
+              "unique_id": lambda: "goodbye",
+          },
+      )
+    self.assertIn(
+        "An implementation for function overload id 'unique_id' already"
+        " exists.",
+        str(e.exception),
+    )
+
 
 class TestCelExtension(cel.CelExtension):
   """An example CEL extension for testing."""

cel_expr_python/py_cel_env.cc

@@ -31,6 +31,7 @@
 #include "cel_expr_python/py_cel_env_config.h"
 #include "cel_expr_python/py_cel_env_internal.h"
 #include "cel_expr_python/py_cel_expression.h"
+#include "cel_expr_python/py_cel_function_decl.h"
 #include "cel_expr_python/py_cel_type.h"
 #include "cel_expr_python/py_error_status.h"
 #include <pybind11/pybind11.h>
@@ -44,10 +45,14 @@ void PyCelEnv::DefinePythonBindings(pybind11::module& m) {
   py::class_<PyCelEnv, std::shared_ptr<PyCelEnv>> cel_class(m, "Env");
   m.def(
       "NewEnv",
-      [](py::object descriptor_pool, std::optional<PyCelEnvConfig> config,
-         std::optional<std::unordered_map<std::string, PyCelType>> variables,
-         std::optional<std::vector<py::object>> extensions,
-         const std::optional<std::string>& container) {
+      [](py::object descriptor_pool, std::optional<PyCelEnvConfig>& config,
+         std::optional<std::unordered_map<std::string, PyCelType>>& variables,
+         std::optional<std::vector<py::object>>& extensions,
+         const std::optional<std::string>& container,
+         std::optional<std::vector<std::shared_ptr<PyCelFunctionDecl>>>&
+             functions,
+         std::optional<std::unordered_map<std::string, py::object>>&
+             function_impls) {
         PyObject* pool_ptr;
         if (descriptor_pool.is_none()) {
           // Replicates python's `descriptor_pool.Default()`
@@ -75,11 +80,16 @@ void PyCelEnv::DefinePythonBindings(pybind11::module& m) {
         return PyCelEnv(config.value_or(PyCelEnvConfig()), pool_ptr,
                         std::move(variables).value_or(
                             std::unordered_map<std::string, PyCelType>{}),
-                        ext_ptrs, container.value_or(""));
+                        ext_ptrs, container.value_or(""),
+                        functions.value_or(
+                            std::vector<std::shared_ptr<PyCelFunctionDecl>>{}),
+                        function_impls.value_or(
+                            std::unordered_map<std::string, py::object>{}));
       },
       py::arg("descriptor_pool") = py::none(), py::arg("config") = py::none(),
       py::arg("variables") = py::none(), py::arg("extensions") = py::none(),
-      py::arg("container") = py::none());
+      py::arg("container") = py::none(), py::arg("functions") = py::none(),
+      py::arg("function_impls") = py::none());
   cel_class
       .def("config",
            [](PyCelEnv& self) { return self.GetEnv()->GetEnvConfig(); })
@@ -112,13 +122,15 @@ void PyCelEnv::DefinePythonBindings(pybind11::module& m) {
           py::arg("arena") = nullptr);
 }
 
-PyCelEnv::PyCelEnv(const PyCelEnvConfig& config, PyObject* descriptor_pool,
-                   std::unordered_map<std::string, PyCelType> variable_types,
-                   const std::vector<PyObject*>& extensions,
-                   std::string container) {
+PyCelEnv::PyCelEnv(
+    const PyCelEnvConfig& config, PyObject* descriptor_pool,
+    const std::unordered_map<std::string, PyCelType>& variable_types,
+    const std::vector<PyObject*>& extensions, const std::string& container,
+    const std::vector<std::shared_ptr<PyCelFunctionDecl>>& functions,
+    const std::unordered_map<std::string, py::object>& function_impls) {
   env_ = ThrowIfError(PyCelEnvInternal::NewCelEnvInternal(
       config, descriptor_pool, std::move(variable_types), extensions,
-      std::move(container)));
+      std::move(container), std::move(functions), std::move(function_impls)));
   ABSL_CHECK(PyGILState_Check());
 }
 

cel_expr_python/py_cel_env.h

@@ -28,6 +28,7 @@
 #include "cel_expr_python/py_cel_env_config.h"
 #include "cel_expr_python/py_cel_expression.h"
 #include "cel_expr_python/py_cel_function.h"
+#include "cel_expr_python/py_cel_function_decl.h"
 #include "cel_expr_python/py_cel_type.h"
 #include <pybind11/pybind11.h>
 
@@ -67,8 +68,11 @@ class PyCelEnv {
  private:
   // Private constructor. Use `py_cel.NewEnv()` in python to obtain an instance.
   PyCelEnv(const PyCelEnvConfig& config, PyObject* descriptor_pool,
-           std::unordered_map<std::string, PyCelType> variable_types,
-           const std::vector<PyObject*>& extensions, std::string container);
+           const std::unordered_map<std::string, PyCelType>& variable_types,
+           const std::vector<PyObject*>& extensions,
+           const std::string& container,
+           const std::vector<std::shared_ptr<PyCelFunctionDecl>>& functions,
+           const std::unordered_map<std::string, py::object>& function_impls);
 
   std::shared_ptr<PyCelEnvInternal> env_;
 };