Add extension version support

PiperOrigin-RevId: 893216967

Author: dmitriplotnikov

MODULE.bazel

@@ -14,7 +14,7 @@ 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
+# 04/01/2026
 _CEL_CPP_COMMIT = "5a3463337cf2a9b90b53833af2bbc1f35da90d64"
 
 _CEL_CPP_SHA256 = "299be398d1495340eb92da31f9a0667e1351479752e8a567ac31c385e4aea73c"

README.md

@@ -234,36 +234,34 @@ expr = cel_env.compile("my_func(1)")
 
 To define a custom extension in C++, define a class extending
 `cel_python::CelExtension`. There are two methods you will need to implement:
-`ConfigureCompiler` and `ConfigureRuntime`. The implementations of these methods
-use the same API as extensions written for the C++ CEL runtime. In fact,
+`GetCompilerLibrary` and `ConfigureRuntime`. The implementations of these
+methods use the same API as extensions written for the C++ CEL runtime. In fact,
 extensions written for the C++ runtime can be used unchanged with
 cel-expr-python - you would just need to write a trivial wrapper class invoking
 the registration functions defined by the C++ extension.
 
-```cpp
-  absl::Status ConfigureCompiler(
-      cel::CompilerBuilder& compiler_builder,
-      const proto2::DescriptorPool& descriptor_pool);
-```
 This method adds extension function definitions to the provided
 `CompilerBuilder`, for example:
 
 ```cpp
-absl::Status ConfigureCompiler(
-      cel::CompilerBuilder& compiler_builder,
-      const proto2::DescriptorPool& descriptor_pool) {
-    CEL_PYTHON_ASSIGN_OR_RETURN(
-        auto func_translate,
-        cel::MakeFunctionDecl("translate",
-            cel::MakeMemberOverloadDecl("translate_inst",
-                                /*return_type=*/cel::StringType(),
-                                /*target=*/cel::StringType(),
-                                /*from_lang=*/cel::StringType(),
-                                /*to_lang=*/cel::StringType())));
-    CEL_PYTHON_RETURN_IF_ERROR(
-        compiler_builder.GetCheckerBuilder().AddFunction(func_translate));
-    return absl::OkStatus();
-}
+  cel::CompilerLibrary GetCompilerLibrary() {
+    return cel::CompilerLibrary(
+        "translate-ext",
+        [](cel::TypeCheckerBuilder& checker_builder) -> absl::Status {
+          CEL_PYTHON_ASSIGN_OR_RETURN(
+              auto func_translate,
+              cel::MakeFunctionDecl(
+                  "translate",
+                  cel::MakeMemberOverloadDecl("translate_inst",
+                                              /*return_type=*/cel::StringType(),
+                                              /*target=*/cel::StringType(),
+                                              /*from_lang=*/cel::StringType(),
+                                              /*to_lang=*/cel::StringType())));
+          CEL_PYTHON_RETURN_IF_ERROR(
+              checker_builder.AddFunction(func_translate));
+          return absl::OkStatus();
+        });
+  }
 ```
 
 The other method registers the actual implementation

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",
@@ -99,14 +104,25 @@ pybind_extension(
 # For pybind11-based CEL extensions.
 pybind_library(
     name = "cel_extension",
-    hdrs = ["cel_extension.h"],
+    srcs = [
+        "py_error_status.cc",
+    ],
+    hdrs = [
+        "cel_extension.h",
+        "py_error_status.h",
+    ],
     visibility = ["//visibility:public"],
     deps = [
+        ":status_macros",
+        "@com_google_absl//absl/base:no_destructor",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/log:absl_check",
+        "@com_google_absl//absl/log:absl_log",
         "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
         "@com_google_cel_cpp//compiler",
         "@com_google_cel_cpp//runtime:runtime_builder",
         "@com_google_cel_cpp//runtime:runtime_options",
-        "@com_google_protobuf//:protobuf",
     ],
 )
 
@@ -135,6 +151,10 @@ py_test(
     srcs = ["cel_env_test.py"],
     deps = [
         ":cel",
+        "//cel_expr_python/ext:ext_bindings",
+        "//cel_expr_python/ext:ext_math",
+        "//cel_expr_python/ext:ext_optional",
+        "//cel_expr_python/ext:ext_strings",
         "//testing:proto2_test_all_types_py_pb2",
         "@com_google_absl_py//absl/testing:absltest",
     ],

cel_expr_python/cel_env_test.py

@@ -18,8 +18,14 @@
 ability to be created from and serialized to YAML format.
 """
 
+import textwrap
+
 from absl.testing import absltest
 from cel_expr_python import cel
+from cel_expr_python.ext import ext_bindings
+from cel_expr_python.ext import ext_math
+from cel_expr_python.ext import ext_optional
+from cel_expr_python.ext import ext_strings
 from cel.expr.conformance.proto2 import test_all_types_pb2 as test_all_types_pb
 
 
@@ -91,6 +97,16 @@ def test_invalid_yaml(self):
         str(e.exception),
     )
 
+  def test_config_export_container(self):
+    env = cel.NewEnv(container="test.container")
+    yaml = env.config().to_yaml()
+    self.assertEqual(
+        normalize_yaml(yaml),
+        normalize_yaml("""
+          container: "test.container"
+        """),
+    )
+
   def test_config_export_variables(self):
     config = cel.NewEnv(
         variables={
@@ -236,27 +252,148 @@ def test_config_variable_types(self):
     self.assertEqual(res.type(), cel.Type.INT)
     self.assertEqual(res.value(), 42)
 
+  def test_config_export_extension_version(self):
+    env = cel.NewEnv(
+        extensions=[
+            ext_math.ExtMath(0),
+            ext_optional.ExtOptional(1),
+            ext_strings.ExtStrings(2),
+            ext_bindings.ExtBindings(),
+        ],
+    )
+    yaml = env.config().to_yaml()
+    self.assertEqual(
+        normalize_yaml(yaml),
+        normalize_yaml("""
+          extensions:
+            - name: "bindings"
+            - name: "math"
+              version: 0
+            - name: "optional"
+              version: 1
+            - name: "strings"
+              version: 2
+        """),
+    )
+
+  def test_config_extension_version_out_of_range(self):
+    cases = [
+        [
+            lambda: ext_math.ExtMath(42),
+            r"'math' extension version: 42 not in range \[0, \d+\]",
+        ],
+        [
+            lambda: ext_optional.ExtOptional(6),
+            r"'optional' extension version: 6 not in range \[0, \d+\]",
+        ],
+        [
+            lambda: ext_strings.ExtStrings(18),
+            r"'strings' extension version: 18 not in range \[0, \d+\]",
+        ],
+    ]
+    for test_case in cases:
+      with self.assertRaises(Exception) as e:
+        cel.NewEnv(
+            extensions=[test_case[0]()],
+        )
+      self.assertRegex(str(e.exception), test_case[1])
+
+  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_extension_override_same_version(self):
+    config = cel.NewEnvConfigFromYaml("""
+      extensions:
+        - name: cel.lib.ext.math
+          version: 1
+        - name: strings
+          version: 2
+      """)
+    env = cel.NewEnv(
+        config=config,
+        extensions=[ext_math.ExtMath(1), ext_strings.ExtStrings(2)],
+    )
+    res = env.compile("'%.3f'.format([math.floor(3.14)])").eval()
+    self.assertEqual(res.value(), "3.000")
+
+  def test_config_extension_override_different_version(self):
+    config = cel.NewEnvConfigFromYaml("""
+      extensions:
+        - name: math
+          version: 0
+        - name: cel.lib.ext.strings
+          version: 2
+      """)
+    with self.assertRaises(Exception) as e:
+      cel.NewEnv(
+          config=config,
+          extensions=[ext_math.ExtMath()],
+      )
+    self.assertIn(
+        "Extension 'math' version 0 is already included. Cannot"
+        " also include version 2",
+        str(e.exception),
+    )
+    with self.assertRaises(Exception) as e:
+      cel.NewEnv(
+          config=config,
+          extensions=[ext_strings.ExtStrings(1)],
+      )
+    self.assertIn(
+        "Extension 'cel.lib.ext.strings' version 2 is already included. Cannot"
+        " also include version 1",
+        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")
-  indent = -1
-  unindented_lines = []
-  for line in lines:
-    pos = -1
-    for i, char in enumerate(line):
-      if char != " " and char != "\t":
-        pos = i
-        break
-    if pos == -1:
-      # Skip blank lines.
-      continue
-    if indent == -1:
-      indent = pos
-    if pos >= indent:
-      unindented_lines.append(line[indent:])
-    else:
-      unindented_lines.append(line)
-  return "\n".join(unindented_lines)
+  return textwrap.dedent(yaml).strip()
 
 
 if __name__ == "__main__":