Make extensions reusable across different environments

PiperOrigin-RevId: 852967383

Author: dmitriplotnikov

BUILD

@@ -81,6 +81,7 @@ pybind_extension(
         "@com_google_cel_cpp//parser:parser_interface",
         "@com_google_cel_cpp//runtime",
         "@com_google_cel_cpp//runtime:activation",
+        "@com_google_cel_cpp//runtime:embedder_context",
         "@com_google_cel_cpp//runtime:function",
         "@com_google_cel_cpp//runtime:reference_resolver",
         "@com_google_cel_cpp//runtime:runtime_builder",

py_cel_activation.cc

@@ -37,6 +37,9 @@ namespace cel_python {
 
 namespace py = ::pybind11;
 
+static const cel::FunctionDescriptorOptions kFunctionDescriptorOptions = {
+    .is_strict = true, .is_contextual = true};
+
 void PyCelActivation::DefinePythonBindings(py::module& m) {
   py::class_<PyCelActivation, std::shared_ptr<PyCelActivation>>(m,
                                                                 "Activation");
@@ -67,10 +70,10 @@ PyCelActivation::PyCelActivation(
     }
     cel::FunctionDescriptor func_descriptor(function->function_name(),
                                             function->is_member(), parameters,
-                                            /*is_strict=*/true);
+                                            kFunctionDescriptorOptions);
     activation_.InsertFunction(
         func_descriptor, std::make_unique<PyCelFunctionAdapter>(
-                             env, function->function_name(), function->impl()));
+                             function->function_name(), function->impl()));
   }
 };
 

py_cel_env.cc

@@ -172,10 +172,7 @@ absl::StatusOr<PyCelExtension*> PyCelExtensionHandle::GetExtension(
   absl::Status status_py_cel_extension;
   try {
     pybind11::handle handle = pybind11::handle(py_extension_);
-    PyCelPythonExtension* py_cel_extension =
-        handle.cast<PyCelPythonExtension*>();
-    status_py_cel_extension = py_cel_extension->SetEnv(env);
-    return py_cel_extension;
+    return handle.cast<PyCelPythonExtension*>();
   } catch (const pybind11::cast_error& e) {
     status_py_cel_extension = absl::InvalidArgumentError(e.what());
   }

py_cel_expression.cc

@@ -38,6 +38,7 @@
 #include "compiler/compiler.h"
 #include "extensions/protobuf/runtime_adapter.h"
 #include "parser/parser_interface.h"
+#include "runtime/embedder_context.h"
 #include "runtime/runtime.h"
 #include "py_cel_activation.h"
 #include "py_cel_arena.h"
@@ -135,10 +136,14 @@ absl::StatusOr<PyCelValue> PyCelExpression::Eval(
   }
   std::shared_ptr<PyCelArena> arena = activation.GetArena();
   std::shared_ptr<PyCelEnv> env = activation.GetEnv();
+  cel::EmbedderContext embedder_context = cel::EmbedderContext::From(&env);
+  cel::EvaluateOptions options;
+  options.message_factory = env->GetMessageFactory();
+  options.embedder_context = &embedder_context;
   CEL_PYTHON_ASSIGN_OR_RETURN(
       cel::Value result,
-      cel_program_->Evaluate(arena->GetArena(), env->GetMessageFactory(),
-                             *activation.GetActivation()));
+      cel_program_->Evaluate(arena->GetArena(), *activation.GetActivation(),
+                             std::move(options)));
   return PyCelValue(result, arena, std::move(env));
 }
 

py_cel_function.cc

@@ -27,6 +27,7 @@
 #include "absl/strings/str_format.h"
 #include "absl/types/span.h"
 #include "common/value.h"
+#include "runtime/embedder_context.h"
 #include "runtime/function.h"
 #include "py_cel_env.h"
 #include "py_cel_type.h"
@@ -41,6 +42,17 @@ namespace cel_python {
 
 namespace py = ::pybind11;
 
+namespace {
+
+static std::shared_ptr<PyCelEnv> GetEnvFromContext(
+    const cel::Function::InvokeContext& context) {
+  ABSL_CHECK(context.embedder_context());  // Crash OK: all call sites are local
+                                           // to the library.
+  return *context.embedder_context()->Get<std::shared_ptr<PyCelEnv>*>();
+}
+
+}  // namespace
+
 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*>(),
@@ -65,12 +77,9 @@ PyCelFunction::~PyCelFunction() {
   PyGILState_Release(gil_state);
 };
 
-PyCelFunctionAdapter::PyCelFunctionAdapter(const std::shared_ptr<PyCelEnv>& env,
-                                           std::string function_name,
+PyCelFunctionAdapter::PyCelFunctionAdapter(std::string function_name,
                                            PyObject* py_function)
-    : env_(env),
-      function_name_(std::move(function_name)),
-      py_function_(py_function) {
+    : function_name_(std::move(function_name)), py_function_(py_function) {
   Py_XINCREF(py_function_);
 }
 
@@ -85,12 +94,13 @@ absl::StatusOr<cel::Value> PyCelFunctionAdapter::Invoke(
     const cel::Function::InvokeContext& context) const {
   ABSL_CHECK(PyGILState_Check());
 
+  std::shared_ptr<PyCelEnv> env = GetEnvFromContext(context);
   PY_CEL_ASSIGN_OR_RETURN(auto py_arena,
                           PyCelArena::FromProtoArena(context.arena()));
   PyObject* py_args = PyTuple_New(args.size());
   for (int i = 0; i < args.size(); ++i) {
     PyTuple_SetItem(py_args, i,
-                    CelValueToPyObject(args[i], env_, py_arena,
+                    CelValueToPyObject(args[i], env, py_arena,
                                        /*plain_value=*/true));
   }
   PyObject* result = PyObject_CallObject(py_function_, py_args);
@@ -105,7 +115,7 @@ absl::StatusOr<cel::Value> PyCelFunctionAdapter::Invoke(
         return absl::StrFormat("Python function '%s'",
                                PyUnicode_AsUTF8(PyObject_Repr(py_function_)));
       },
-      env_, context.arena());
+      env, context.arena());
 };
 
 }  // namespace cel_python

py_cel_function.h

@@ -62,16 +62,14 @@ class PyCelFunction {
 // function.
 class PyCelFunctionAdapter : public cel::Function {
  public:
-  PyCelFunctionAdapter(const std::shared_ptr<PyCelEnv>& env,
-                       std::string function_name, PyObject* py_function);
+  PyCelFunctionAdapter(std::string function_name, PyObject* py_function);
   ~PyCelFunctionAdapter() override;
 
   absl::StatusOr<cel::Value> Invoke(
       absl::Span<const cel::Value> args,
       const cel::Function::InvokeContext& context) const final;
 
  private:
-  std::shared_ptr<PyCelEnv> env_;
   std::string function_name_;
   PyObject* py_function_;
 };

py_cel_python_extension.cc

@@ -46,6 +46,9 @@ namespace cel_python {
 
 namespace py = ::pybind11;
 
+static const cel::FunctionDescriptorOptions kFunctionDescriptorOptions = {
+    .is_strict = true, .is_contextual = true};
+
 void PyCelPythonExtension::DefinePythonBindings(py::module_& m) {
   py::class_<PyCelExtension>(m, "CelExtensionBase")
       .def(py::init<std::string>(), py::arg("name"));
@@ -59,18 +62,6 @@ PyCelPythonExtension::PyCelPythonExtension(
     std::string name, std::vector<PyCelFunctionDecl> functions)
     : PyCelExtension(std::move(name)), functions_(std::move(functions)) {}
 
-// TODO(b/462745713): pass the env to the Invoke method instead of storing it
-// as a member variable and remove this method.
-absl::Status PyCelPythonExtension::SetEnv(
-    const std::shared_ptr<PyCelEnv>& env) {
-  if (env_ != nullptr && env_ != env) {
-    return absl::FailedPreconditionError(
-        "PyCelExtension already has an environment");
-  }
-  env_ = env;
-  return absl::OkStatus();
-}
-
 absl::Status PyCelPythonExtension::ConfigureCompiler(
     cel::CompilerBuilder& compiler_builder,
     const google::protobuf::DescriptorPool& descriptor_pool) {
@@ -115,12 +106,11 @@ absl::Status PyCelPythonExtension::ConfigureRuntime(
       }
 
       cel::FunctionDescriptor descriptor(function.name(), overload.is_member(),
-                                         types,
-                                         /*is_strict=*/true);
+                                         types, kFunctionDescriptorOptions);
       if (overload.py_function()) {
         PY_CEL_RETURN_IF_ERROR(runtime_builder.function_registry().Register(
             descriptor, std::make_unique<PyCelFunctionAdapter>(
-                            env_, function.name(), overload.py_function())));
+                            function.name(), overload.py_function())));
       } else {
         PY_CEL_RETURN_IF_ERROR(
             runtime_builder.function_registry().RegisterLazyFunction(

py_cel_python_extension.h

@@ -15,7 +15,6 @@
 #ifndef THIRD_PARTY_CEL_PYTHON_PY_CEL_PYTHON_EXTENSION_H_
 #define THIRD_PARTY_CEL_PYTHON_PY_CEL_PYTHON_EXTENSION_H_
 
-#include <memory>
 #include <string>
 #include <vector>
 
@@ -30,18 +29,12 @@
 
 namespace cel_python {
 
-class PyCelEnv;
-
 class PyCelPythonExtension : public PyCelExtension {
  public:
   static void DefinePythonBindings(pybind11::module& m);
   PyCelPythonExtension(std::string name,
                        std::vector<PyCelFunctionDecl> functions);
 
-  // TODO(b/462745713): pass the env to the Invoke method instead of storing it
-  // as a member variable.
-  absl::Status SetEnv(const std::shared_ptr<PyCelEnv>& env);
-
  protected:
   absl::Status ConfigureCompiler(
       cel::CompilerBuilder& compiler_builder,
@@ -52,7 +45,6 @@ class PyCelPythonExtension : public PyCelExtension {
 
  private:
   std::vector<PyCelFunctionDecl> functions_;
-  std::shared_ptr<PyCelEnv> env_;
 };
 
 }  // namespace cel_python