From 0786ac2a9edea787353e3d7ead8ffbb3a9f4223e Mon Sep 17 00:00:00 2001 From: Dmitri Plotnikov Date: Tue, 7 Apr 2026 17:23:14 -0700 Subject: [PATCH] Integrate cel::Config::ExtensionConfig into Python CEL environment PiperOrigin-RevId: 896167243 --- MODULE.bazel | 6 +- cel_expr_python/BUILD | 9 +- cel_expr_python/cel_env_test.py | 70 ++++++++++ cel_expr_python/ext/ext_optional.cc | 2 +- cel_expr_python/py_cel_env_internal.cc | 182 ++++++++++++++++++++----- cel_expr_python/py_cel_env_internal.h | 18 ++- 6 files changed, 247 insertions(+), 40 deletions(-) diff --git a/MODULE.bazel b/MODULE.bazel index 51589db..fb0bdd0 100644 --- a/MODULE.bazel +++ b/MODULE.bazel @@ -14,10 +14,10 @@ bazel_dep(name = "bazel_skylib", version = "1.8.2") # https://registry.bazel.build/modules/cel-cpp bazel_dep(name = "cel-cpp", version = "0.14.0", repo_name = "com_google_cel_cpp") -# 03/30/2026 -_CEL_CPP_COMMIT = "5a3463337cf2a9b90b53833af2bbc1f35da90d64" +# 04/07/2025 +_CEL_CPP_COMMIT = "7022451780867bd9c173956f5793b5dfb600928e" -_CEL_CPP_SHA256 = "299be398d1495340eb92da31f9a0667e1351479752e8a567ac31c385e4aea73c" +_CEL_CPP_SHA256 = "d1b9b2c2e9745826c1c8540a49f972ea4b93dffd5878c6a67756a2627b0119ef" archive_override( module_name = "cel-cpp", diff --git a/cel_expr_python/BUILD b/cel_expr_python/BUILD index 35624d6..d4b206a 100644 --- 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,7 @@ 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_cpp//validator", "@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 +141,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 index 5b49875..cd26c77 100644 --- 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 index 80a01e2..72b7086 100644 --- 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 index c92820b..3c50900 100644 --- a/cel_expr_python/py_cel_env_internal.cc +++ b/cel_expr_python/py_cel_env_internal.cc @@ -25,30 +25,120 @@ #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 "validator/validator.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 namespace cel_python { -PyCelEnvInternal::PyCelEnvInternal(const PyCelEnvConfig& env_config, - PyObject* py_descriptor_pool, - const std::vector& extensions, - const std::string& container) +namespace { + +// 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_; + } + + cel::Validator& GetValidator() override { return validator_; } + + absl::StatusOr> Build() override { + return absl::UnimplementedError("Not implemented"); + } + + private: + cel::ParserBuilder* parser_builder_ = nullptr; + cel::TypeCheckerBuilder* checker_builder_ = nullptr; + cel::Validator validator_; +}; + +// A temporary CompilerLibrary that deals with the interface mismatch between +// CelExtension and cel::CompilerLibrary. +cel::CompilerLibrary MakeCompilerLibrary( + CelExtension* extension, + const std::shared_ptr& descriptor_pool, + cel::ParserBuilder* passive_parser_builder, + cel::TypeCheckerBuilder* passive_checker_builder) { + return 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 { + CompilerBuilderAdapter compiler_builder(&parser_builder, + passive_checker_builder); + return extension->ConfigureCompiler(compiler_builder, *descriptor_pool); + }, + // Safe to capture passive_parser_builder because it outlives the + // compiler library. + [extension, descriptor_pool, passive_parser_builder]( + cel::TypeCheckerBuilder& checker_builder) -> absl::Status { + CompilerBuilderAdapter compiler_builder(passive_parser_builder, + &checker_builder); + return extension->ConfigureCompiler(compiler_builder, *descriptor_pool); + }); +} +} // namespace + +PyCelEnvInternal::PyCelEnvInternal( + const PyCelEnvConfig& env_config, PyObject* py_descriptor_pool, + std::vector extension_handles, + const std::string& container) : env_config_(env_config), py_descriptor_database_(py_descriptor_pool), descriptor_pool_( @@ -56,11 +146,36 @@ PyCelEnvInternal::PyCelEnvInternal(const PyCelEnvConfig& env_config, message_factory_(descriptor_pool_.get()), py_message_factory_( std::make_shared(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(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 (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 MakeCompilerLibrary(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 +194,20 @@ PyCelEnvInternal::NewCelEnvInternal( })); } - return std::shared_ptr(new PyCelEnvInternal( - PyCelEnvConfig(config), py_descriptor_pool, extensions, container)); + std::vector extension_handles; + extension_handles.reserve(extensions.size()); + for (PyObject* ext : extensions) { + CelExtensionHandle handle(ext); + CEL_PYTHON_ASSIGN_OR_RETURN(CelExtension * extension, + handle.GetExtension()); + // TODO(b/498655870): support extension version. + CEL_PYTHON_RETURN_IF_ERROR(config.AddExtensionConfig(extension->name())); + extension_handles.push_back(std::move(handle)); + } + + return std::shared_ptr( + new PyCelEnvInternal(PyCelEnvConfig(config), py_descriptor_pool, + std::move(extension_handles), container)); } absl::StatusOr PyCelEnvInternal::GetCompiler( @@ -95,13 +222,6 @@ absl::StatusOr PyCelEnvInternal::GetCompiler( std::unique_ptr compiler_builder, env->cel_env_.NewCompilerBuilder()); compiler_builder->GetCheckerBuilder().set_container(env->container_); - for (std::unique_ptr& 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 +245,7 @@ absl::StatusOr 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 +257,10 @@ absl::StatusOr 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& 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 runtime, std::move(builder).Build()); const cel::Runtime* runtime_ptr = runtime.get(); @@ -171,14 +284,21 @@ CelExtensionHandle::CelExtensionHandle(PyObject* extension) Py_INCREF(py_extension_); } +CelExtensionHandle::CelExtensionHandle(CelExtensionHandle&& other) + : py_extension_(other.py_extension_), cel_extension_(other.cel_extension_) { + other.py_extension_ = nullptr; + other.cel_extension_ = nullptr; +} + CelExtensionHandle::~CelExtensionHandle() { - auto gil_state = PyGILState_Ensure(); - Py_DECREF(py_extension_); - PyGILState_Release(gil_state); + if (py_extension_ != nullptr) { + auto gil_state = PyGILState_Ensure(); + Py_DECREF(py_extension_); + PyGILState_Release(gil_state); + } } -absl::StatusOr CelExtensionHandle::GetExtension( - const std::shared_ptr& env) { +absl::StatusOr 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 index d6db52a..d9825ab 100644 --- a/cel_expr_python/py_cel_env_internal.h +++ b/cel_expr_python/py_cel_env_internal.h @@ -25,8 +25,11 @@ #include "absl/container/flat_hash_map.h" #include "absl/status/status.h" #include "absl/status/statusor.h" +#include "checker/type_checker_builder.h" #include "compiler/compiler.h" #include "env/env.h" +#include "env/env_runtime.h" +#include "parser/parser_interface.h" #include "runtime/runtime.h" #include "runtime/runtime_builder.h" #include "runtime/runtime_options.h" @@ -46,10 +49,11 @@ class PyCelEnvInternal; class CelExtensionHandle { public: explicit CelExtensionHandle(PyObject* extension); + CelExtensionHandle(const CelExtensionHandle& other) = delete; + CelExtensionHandle(CelExtensionHandle&& other); ~CelExtensionHandle(); - absl::StatusOr GetExtension( - const std::shared_ptr& env); + absl::StatusOr GetExtension(); private: // The Python object that was passed to the constructor and is retained for @@ -105,7 +109,7 @@ class PyCelEnvInternal { // Use NewCelEnvInternal() to create an instance. PyCelEnvInternal(const PyCelEnvConfig& env_config, PyObject* py_descriptor_pool, - const std::vector& extensions, + std::vector extensions, const std::string& container); absl::Status ConfigureStandardExtension( @@ -116,6 +120,7 @@ class PyCelEnvInternal { const cel::RuntimeOptions& opts); cel::Env cel_env_; + cel::EnvRuntime cel_env_runtime_; PyCelEnvConfig env_config_; PyDescriptorDatabase py_descriptor_database_; std::shared_ptr descriptor_pool_; @@ -123,11 +128,16 @@ class PyCelEnvInternal { std::shared_ptr py_message_factory_; // Synchronized by the GIL. std::unordered_map variable_types_; - std::vector> extensions_; + std::vector extensions_; std::string container_; std::unique_ptr compiler_; absl::flat_hash_map> runtimes_; + + // Passive parser and checker builders function as sinks to configure custom + // extensions, but are themselves not used to build a compiler. + std::unique_ptr passive_parser_builder_; + std::unique_ptr passive_checker_builder_; }; } // namespace cel_python