Integrate cel::Config::ExtensionConfig into Python CEL environment

dmitriplotnikov · copybara-github · commit 4eb8e787be9e · 2026-04-01T14:33:43.000-07:00
PiperOrigin-RevId: 892660001
diff --git a/cel_expr_python/BUILD b/cel_expr_python/BUILD
@@ -63,6 +63,7 @@ pybind_extension(
         "@com_google_absl//absl/types:optional",
         "@com_google_absl//absl/types:span",
         "@com_google_cel_cpp//checker:type_checker_builder",
+        "@com_google_cel_cpp//checker:type_checker_builder_factory",
         "@com_google_cel_cpp//checker:validation_result",
         "@com_google_cel_cpp//common:ast",
         "@com_google_cel_cpp//common:ast_proto",
@@ -78,8 +79,13 @@ pybind_extension(
         "@com_google_cel_cpp//compiler",
         "@com_google_cel_cpp//env",
         "@com_google_cel_cpp//env:config",
+        "@com_google_cel_cpp//env:env_runtime",
+        "@com_google_cel_cpp//env:env_std_extensions",
         "@com_google_cel_cpp//env:env_yaml",
+        "@com_google_cel_cpp//env:runtime_std_extensions",
         "@com_google_cel_cpp//extensions/protobuf:runtime_adapter",
+        "@com_google_cel_cpp//parser",
+        "@com_google_cel_cpp//parser:options",
         "@com_google_cel_cpp//parser:parser_interface",
         "@com_google_cel_cpp//runtime",
         "@com_google_cel_cpp//runtime:activation",
@@ -88,7 +94,6 @@ pybind_extension(
         "@com_google_cel_cpp//runtime:reference_resolver",
         "@com_google_cel_cpp//runtime:runtime_builder",
         "@com_google_cel_cpp//runtime:runtime_options",
-        "@com_google_cel_cpp//runtime:standard_runtime_builder_factory",
         "@com_google_cel_spec//proto/cel/expr:checked_cc_proto",
         "@com_google_cel_spec//proto/cel/expr:syntax_cc_proto",
         "@com_google_protobuf//:protobuf",
@@ -135,6 +140,7 @@ py_test(
     srcs = ["cel_env_test.py"],
     deps = [
         ":cel",
+        "//cel_expr_python/ext:ext_math",
         "//testing:proto2_test_all_types_py_pb2",
         "@com_google_absl_py//absl/testing:absltest",
     ],
diff --git a/cel_expr_python/cel_env_test.py b/cel_expr_python/cel_env_test.py
@@ -20,6 +20,7 @@
 
 from absl.testing import absltest
 from cel_expr_python import cel
+from cel_expr_python.ext import ext_math
 from cel.expr.conformance.proto2 import test_all_types_pb2 as test_all_types_pb
 
 
@@ -236,6 +237,75 @@ def test_config_variable_types(self):
     self.assertEqual(res.type(), cel.Type.INT)
     self.assertEqual(res.value(), 42)
 
+  def test_config_extensions(self):
+    config = cel.NewEnvConfigFromYaml("""
+      extensions:
+        - name: math
+        - name: strings
+      """)
+    env = cel.NewEnv(
+        config=config,
+        extensions=[TestCelExtension()],
+    )
+    yaml = env.config().to_yaml()
+    self.assertEqual(
+        normalize_yaml(yaml),
+        normalize_yaml("""
+          extensions:
+            - name: "math"
+            - name: "strings"
+            - name: "test_cel_extension"
+        """),
+    )
+    res = env.compile("'%.4f'.format([math.sqrt(2)])").eval()
+    self.assertEqual(res.value(), "1.4142")
+    res = env.compile("hello('World')").eval()
+    self.assertEqual(res.value(), "Hello, World!")
+
+  def test_config_extensions_override(self):
+    # TODO(b/498655870): add assertion based on extension aliases once
+    # supported.
+    config = cel.NewEnvConfigFromYaml("""
+      extensions:
+        - name: cel.lib.ext.math
+          version: 0
+        - name: cel.lib.ext.strings
+      """)
+    with self.assertRaises(Exception) as e:
+      cel.NewEnv(
+          config=config,
+          extensions=[ext_math.ExtMath()],
+      )
+    self.assertIn(
+        "Extension 'cel.lib.ext.math' version 0 is already included. Cannot"
+        " also include version 'latest'",
+        str(e.exception),
+    )
+
+
+class TestCelExtension(cel.CelExtension):
+  """An example CEL extension for testing."""
+
+  def __init__(self):
+    super().__init__(
+        "test_cel_extension",
+        functions=[
+            cel.FunctionDecl(
+                "hello",
+                [
+                    cel.Overload(
+                        "hello(string)",
+                        return_type=cel.Type.STRING,
+                        parameters=[
+                            cel.Type.STRING,
+                        ],
+                        impl=lambda arg: f"Hello, {arg}!",
+                    )
+                ],
+            ),
+        ],
+    )
+
 
 def normalize_yaml(yaml: str) -> str:
   lines = yaml.split("\n")
diff --git a/cel_expr_python/ext/ext_optional.cc b/cel_expr_python/ext/ext_optional.cc
@@ -25,7 +25,7 @@ namespace cel_python {
 
 class ExtOptional : public CelExtension {
  public:
-  explicit ExtOptional() : CelExtension("cel.lib.optional") {}
+  explicit ExtOptional() : CelExtension("optional") {}
 
   absl::Status ConfigureCompiler(
       cel::CompilerBuilder& compiler_builder,
diff --git a/cel_expr_python/py_cel_env_internal.cc b/cel_expr_python/py_cel_env_internal.cc
@@ -25,42 +25,169 @@
 #include "absl/status/statusor.h"
 #include "absl/strings/str_cat.h"
 #include "checker/type_checker_builder.h"
+#include "checker/type_checker_builder_factory.h"
+#include "common/type.h"
 #include "compiler/compiler.h"
 #include "env/config.h"
 #include "env/env.h"
+#include "env/env_std_extensions.h"
+#include "env/runtime_std_extensions.h"
 #include "env/type_info.h"
+#include "parser/options.h"
+#include "parser/parser.h"
+#include "parser/parser_interface.h"
 #include "runtime/reference_resolver.h"
 #include "runtime/runtime.h"
 #include "runtime/runtime_builder.h"
 #include "runtime/runtime_options.h"
-#include "runtime/standard_runtime_builder_factory.h"
 #include "cel_expr_python/cel_extension.h"
 #include "cel_expr_python/py_cel_env_config.h"
 #include "cel_expr_python/py_cel_python_extension.h"
 #include "cel_expr_python/py_cel_type.h"
 #include "cel_expr_python/py_descriptor_database.h"
+#include "cel_expr_python/py_error_status.h"
+#include "cel_expr_python/py_message_factory.h"
 #include "cel_expr_python/status_macros.h"
 #include "google/protobuf/arena.h"
+#include "google/protobuf/descriptor.h"
 #include <pybind11/pybind11.h>
 
 namespace cel_python {
 
-PyCelEnvInternal::PyCelEnvInternal(const PyCelEnvConfig& env_config,
-                                   PyObject* py_descriptor_pool,
-                                   const std::vector<PyObject*>& extensions,
-                                   const std::string& container)
+// A temporary adapter for cel::CompilerBuilder.  It will be removed once the
+// CelExtension class is changed to return a cel::CompilerLibrary directly.
+//
+// This adapter allows `CelExtension::ConfigureCompiler`, which takes a
+// `cel::CompilerBuilder` to be used with the `cel::Env` API, which
+// is expressed in terms of the `cel::CompilerLibrary` framework.  The adapter
+// splits the `cel::CompilerLibrary` into its constituent parts and feeds them
+// to the `cel::CompilerBuilder` interface.
+class CompilerBuilderAdapter : public cel::CompilerBuilder {
+ public:
+  CompilerBuilderAdapter(cel::ParserBuilder* parser_builder,
+                         cel::TypeCheckerBuilder* checker_builder)
+      : parser_builder_(parser_builder), checker_builder_(checker_builder) {}
+
+  absl::Status AddLibrary(cel::CompilerLibrary library) override {
+    if (library.configure_parser != nullptr) {
+      CEL_PYTHON_RETURN_IF_ERROR(library.configure_parser(*parser_builder_));
+    }
+    if (library.configure_checker != nullptr) {
+      CEL_PYTHON_RETURN_IF_ERROR(library.configure_checker(*checker_builder_));
+    }
+    return absl::OkStatus();
+  }
+
+  absl::Status AddLibrarySubset(cel::CompilerLibrarySubset subset) override {
+    return absl::UnimplementedError("Not implemented");
+  }
+
+  cel::ParserBuilder& GetParserBuilder() override {
+    ABSL_CHECK(parser_builder_ != nullptr);
+    return *parser_builder_;
+  }
+
+  cel::TypeCheckerBuilder& GetCheckerBuilder() override {
+    ABSL_CHECK(checker_builder_ != nullptr);
+    return *checker_builder_;
+  }
+
+  absl::StatusOr<std::unique_ptr<cel::Compiler>> Build() override {
+    return absl::UnimplementedError("Not implemented");
+  }
+
+ private:
+  cel::ParserBuilder* parser_builder_;
+  cel::TypeCheckerBuilder* checker_builder_;
+};
+
+// A temporary class to deal with the interface mismatch between CelExtension
+// and cel::CompilerLibrary.
+class AdapterCompilerLibrary : public cel::CompilerLibrary {
+ public:
+  AdapterCompilerLibrary(
+      CelExtension* extension,
+      const std::shared_ptr<google::protobuf::DescriptorPool>& descriptor_pool,
+      cel::ParserBuilder* passive_parser_builder,
+      cel::TypeCheckerBuilder* passive_checker_builder)
+      : cel::CompilerLibrary(
+            extension->name(),
+            // Safe to capture passive_checker_builder because it outlives the
+            // compiler library.
+            [extension, descriptor_pool, passive_checker_builder](
+                cel::ParserBuilder& parser_builder) -> absl::Status {
+              return ConfigureParser(parser_builder, extension, descriptor_pool,
+                                     passive_checker_builder);
+            },
+            // Safe to capture passive_parser_builder because it outlives the
+            // compiler library.
+            [extension, descriptor_pool, passive_parser_builder](
+                cel::TypeCheckerBuilder& checker_builder) -> absl::Status {
+              return ConfigureChecker(checker_builder, extension,
+                                      descriptor_pool, passive_parser_builder);
+            }) {};
+
+ private:
+  static absl::Status ConfigureParser(
+      cel::ParserBuilder& parser_builder, CelExtension* extension,
+      const std::shared_ptr<google::protobuf::DescriptorPool>& descriptor_pool,
+      cel::TypeCheckerBuilder* passive_checker_builder) {
+    CompilerBuilderAdapter compiler_builder(&parser_builder,
+                                            passive_checker_builder);
+    return extension->ConfigureCompiler(compiler_builder, *descriptor_pool);
+  }
+
+  static absl::Status ConfigureChecker(
+      cel::TypeCheckerBuilder& checker_builder, CelExtension* extension,
+      const std::shared_ptr<google::protobuf::DescriptorPool>& descriptor_pool,
+      cel::ParserBuilder* passive_parser_builder) {
+    CompilerBuilderAdapter compiler_builder(passive_parser_builder,
+                                            &checker_builder);
+    return extension->ConfigureCompiler(compiler_builder, *descriptor_pool);
+  }
+};
+
+PyCelEnvInternal::PyCelEnvInternal(
+    const PyCelEnvConfig& env_config, PyObject* py_descriptor_pool,
+    std::vector<std::unique_ptr<CelExtensionHandle>>& extension_handles,
+    const std::string& container)
     : env_config_(env_config),
       py_descriptor_database_(py_descriptor_pool),
       descriptor_pool_(
           std::make_shared<google::protobuf::DescriptorPool>(&py_descriptor_database_)),
       message_factory_(descriptor_pool_.get()),
       py_message_factory_(
           std::make_shared<PyMessageFactory>(py_descriptor_pool)),
+      extensions_(std::move(extension_handles)),
       container_(std::move(container)) {
   cel_env_.SetDescriptorPool(descriptor_pool_);
   cel_env_.SetConfig(env_config_.GetConfig());
-  for (PyObject* ext : extensions) {
-    extensions_.push_back(std::make_unique<CelExtensionHandle>(ext));
+  cel::RegisterStandardExtensions(cel_env_);
+
+  cel_env_runtime_.SetDescriptorPool(descriptor_pool_);
+  cel_env_runtime_.SetConfig(env_config_.GetConfig());
+  cel::RegisterStandardExtensions(cel_env_runtime_);
+
+  passive_parser_builder_ = cel::NewParserBuilder(cel::ParserOptions());
+  passive_checker_builder_ =
+      ThrowIfError(cel::CreateTypeCheckerBuilder(descriptor_pool_.get()));
+  for (std::unique_ptr<CelExtensionHandle>& extension_handle : extensions_) {
+    // This should never fail because we have already called GetExtension() once
+    // before calling this constructor.
+    CelExtension* extension = ThrowIfError(extension_handle->GetExtension());
+    cel_env_.RegisterCompilerLibrary(
+        extension->name(), extension->name(), 0, [this, extension]() {
+          return AdapterCompilerLibrary(extension, descriptor_pool_,
+                                        passive_parser_builder_.get(),
+                                        passive_checker_builder_.get());
+        });
+    cel_env_runtime_.RegisterExtensionFunctions(
+        extension->name(), extension->name(), 0,
+        [extension](
+            cel::RuntimeBuilder& runtime_builder,
+            const cel::RuntimeOptions& runtime_options) -> absl::Status {
+          return extension->ConfigureRuntime(runtime_builder, runtime_options);
+        });
   }
 }
 
@@ -79,8 +206,19 @@ PyCelEnvInternal::NewCelEnvInternal(
         }));
   }
 
-  return std::shared_ptr<PyCelEnvInternal>(new PyCelEnvInternal(
-      PyCelEnvConfig(config), py_descriptor_pool, extensions, container));
+  std::vector<std::unique_ptr<CelExtensionHandle>> extension_handles;
+  for (PyObject* ext : extensions) {
+    auto extension_handle = std::make_unique<CelExtensionHandle>(ext);
+    CEL_PYTHON_ASSIGN_OR_RETURN(CelExtension * extension,
+                                extension_handle->GetExtension());
+    // TODO(b/498655870): support extension version.
+    CEL_PYTHON_RETURN_IF_ERROR(config.AddExtensionConfig(extension->name()));
+    extension_handles.push_back(std::move(extension_handle));
+  }
+
+  return std::shared_ptr<PyCelEnvInternal>(
+      new PyCelEnvInternal(PyCelEnvConfig(config), py_descriptor_pool,
+                           extension_handles, container));
 }
 
 absl::StatusOr<const cel::Compiler*> PyCelEnvInternal::GetCompiler(
@@ -95,13 +233,6 @@ absl::StatusOr<const cel::Compiler*> PyCelEnvInternal::GetCompiler(
       std::unique_ptr<cel::CompilerBuilder> compiler_builder,
       env->cel_env_.NewCompilerBuilder());
   compiler_builder->GetCheckerBuilder().set_container(env->container_);
-  for (std::unique_ptr<CelExtensionHandle>& extension_handle :
-       env->extensions_) {
-    CEL_PYTHON_ASSIGN_OR_RETURN(CelExtension * extension,
-                                extension_handle->GetExtension(env));
-    CEL_PYTHON_RETURN_IF_ERROR(extension->ConfigureCompiler(
-        *compiler_builder, *(env->descriptor_pool_.get())));
-  }
 
   // Convert variable types from cel::TypeInfo to PyCelType.
   google::protobuf::Arena* arena = compiler_builder->GetCheckerBuilder().arena();
@@ -125,7 +256,7 @@ absl::StatusOr<const cel::Runtime*> PyCelEnvInternal::GetRuntime(
     return it->second.get();
   }
 
-  cel::RuntimeOptions opts;
+  cel::RuntimeOptions& opts = env->cel_env_runtime_.mutable_runtime_options();
   opts.container = env->container_;
   opts.enable_empty_wrapper_null_unboxing = true;
   opts.enable_qualified_type_identifiers = true;
@@ -137,17 +268,10 @@ absl::StatusOr<const cel::Runtime*> PyCelEnvInternal::GetRuntime(
       opts.fail_on_warnings = false;
       break;
   }
-  CEL_PYTHON_ASSIGN_OR_RETURN(
-      cel::RuntimeBuilder builder,
-      cel::CreateStandardRuntimeBuilder(env->descriptor_pool_.get(), opts));
+  CEL_PYTHON_ASSIGN_OR_RETURN(cel::RuntimeBuilder builder,
+                              env->cel_env_runtime_.CreateRuntimeBuilder());
   CEL_PYTHON_RETURN_IF_ERROR(cel::EnableReferenceResolver(
       builder, cel::ReferenceResolverEnabled::kAlways));
-  for (std::unique_ptr<CelExtensionHandle>& extension_handle :
-       env->extensions_) {
-    CEL_PYTHON_ASSIGN_OR_RETURN(CelExtension * extension,
-                                extension_handle->GetExtension(env));
-    CEL_PYTHON_RETURN_IF_ERROR(extension->ConfigureRuntime(builder, opts));
-  }
   CEL_PYTHON_ASSIGN_OR_RETURN(std::unique_ptr<cel::Runtime> runtime,
                               std::move(builder).Build());
   const cel::Runtime* runtime_ptr = runtime.get();
@@ -177,8 +301,7 @@ CelExtensionHandle::~CelExtensionHandle() {
   PyGILState_Release(gil_state);
 }
 
-absl::StatusOr<CelExtension*> CelExtensionHandle::GetExtension(
-    const std::shared_ptr<PyCelEnvInternal>& env) {
+absl::StatusOr<CelExtension*> CelExtensionHandle::GetExtension() {
   if (cel_extension_) {
     return cel_extension_;
   }
diff --git a/cel_expr_python/py_cel_env_internal.h b/cel_expr_python/py_cel_env_internal.h
