Add support for environment options in Python CEL.

This change allows passing an optional dictionary of options to EnvConfig and NewEnv. It exposes an options property on EnvConfig and uses these options to configure the underlying CEL compiler builder, specifically supporting the "enable_pratt_parser" option.

PiperOrigin-RevId: 964845782

Author: dmitriplotnikov

cel_expr_python/BUILD

@@ -20,6 +20,7 @@ pybind_library(
         "py_cel_expression.cc",
         "py_cel_function.cc",
         "py_cel_function_decl.cc",
+        "py_cel_options.cc",
         "py_cel_overload.cc",
         "py_cel_python_extension.cc",
         "py_cel_type.cc",
@@ -36,6 +37,7 @@ pybind_library(
         "py_cel_expression.h",
         "py_cel_function.h",
         "py_cel_function_decl.h",
+        "py_cel_options.h",
         "py_cel_overload.h",
         "py_cel_python_extension.h",
         "py_cel_type.h",

cel_expr_python/cel.pyi

@@ -12,6 +12,10 @@ class CelExtension(CelExtensionBase):
 class CelExtensionBase:
   def __init__(self, name: str) -> None: ...
 
+class Options:
+  enable_pratt_parser: bool
+  def __init__(self, enable_pratt_parser: bool = ...) -> None: ...
+
 class EnvConfig:
   @property
   def context_type(self) -> str: ...
@@ -26,6 +30,7 @@ class Env:
   def compile(self, expression: str, disable_check: bool = ...) -> Expression: ...
   def deserialize(self, serialized: str | bytes) -> Expression: ...
   def config(self) -> EnvConfig: ...
+  def options(self) -> Options: ...
 
 class Expression:
   def eval(self, activation: Activation | None = ..., data: Mapping[str, Any] | None = ..., functions=..., arena: _InternalArena = ...) -> Value: ...
@@ -86,6 +91,15 @@ class _InternalArena:
 
 def Arena() -> _InternalArena: ...
 
-def NewEnv(descriptor_pool: proto_descriptor_pool.DescriptorPool | Any | None = ..., config: EnvConfig | None = ..., variables: Mapping[str, Type] | None = ..., extensions: Sequence[CelExtensionBase] | None = ..., container: str | ExpressionContainer | None = ..., functions: Sequence[FunctionDecl] | None = ..., function_impls: Mapping[str, Callable[..., Any]] | None = ...) -> Env: ...
+def NewEnv(
+    descriptor_pool: proto_descriptor_pool.DescriptorPool | Any | None = ...,
+    config: EnvConfig | None = ...,
+    variables: Mapping[str, Type] | None = ...,
+    extensions: Sequence[CelExtensionBase] | None = ...,
+    container: str | ExpressionContainer | None = ...,
+    functions: Sequence[FunctionDecl] | None = ...,
+    function_impls: Mapping[str, Callable[..., Any]] | None = ...,
+    options: Options | None = ...,
+) -> Env: ...
 
 def NewEnvConfigFromYaml(yaml: str) -> EnvConfig: ...

cel_expr_python/cel_env_test.py

@@ -724,6 +724,20 @@ def test_config_functions_deprecated_syntax(self):
     res = env.compile("'bad'.is_ok()").eval()
     self.assertFalse(res.value())
 
+  def test_env_options(self):
+    options = cel.Options(enable_pratt_parser=True)
+    self.assertTrue(options.enable_pratt_parser)
+    self.assertEqual(repr(options), "Options(enable_pratt_parser=True)")
+    options.enable_pratt_parser = False
+    self.assertFalse(options.enable_pratt_parser)
+    self.assertEqual(repr(options), "Options(enable_pratt_parser=False)")
+
+    env = cel.NewEnv(options=cel.Options(enable_pratt_parser=True))
+    self.assertTrue(env.options().enable_pratt_parser)
+
+    default_env = cel.NewEnv()
+    self.assertFalse(default_env.options().enable_pratt_parser)
+
 
 class TestCelExtension(cel.CelExtension):
   """An example CEL extension for testing."""

cel_expr_python/cel_test.py

@@ -28,7 +28,9 @@
 from cel.expr.conformance.proto2 import test_all_types_pb2 as test_all_types_pb
 
 
-class CelTest(absltest.TestCase):
+@absltest.skipThisClass("Base class")
+class _CelTestBase(absltest.TestCase):
+  options: cel.Options = cel.Options()
 
   def setUp(self):
     super().setUp()
@@ -50,7 +52,8 @@ def setUp(self):
             "var_string_map": cel.Type.Map(cel.Type.STRING, cel.Type.BOOL),
             "var_dyn_map": cel.Type.MAP,
             "var_dyn": cel.Type.DYN,
-        }
+        },
+        options=self.options,
     )
     self.object_counts_before_test = self._grab_object_counts()
 
@@ -615,10 +618,13 @@ def testDynType(self):
     self.assertIn("out of range for 'var_dyn'", res.value())
 
   def testDynType_nonCelType(self):
-    res = self._eval("var_dyn", {"var_dyn": self})
+    class NonCelValue:
+      pass
+
+    res = self._eval("var_dyn", {"var_dyn": NonCelValue()})
     self.assertEqual(res.type(), cel.Type.ERROR)
     self.assertIn(
-        "Non-CEL value type for 'var_dyn': CelTest",
+        "Non-CEL value type for 'var_dyn': NonCelValue",
         res.value(),
     )
 
@@ -768,18 +774,26 @@ def testCompilationErrorHandling(self):
     # Check parser error.
     with self.assertRaises(Exception) as e:
       self.env.compile("'Hello,' # 'World!'", disable_check=True)
-    self.assertIn(
-        "1:10: Syntax error: token recognition error at: '#'\n "
-        "| 'Hello,' # 'World!'\n "
-        "| .........^",
-        str(e.exception),
-    )
-    self.assertIn(
-        "1:12: Syntax error: extraneous input ''World!'' expecting <EOF>\n "
-        "| 'Hello,' # 'World!'\n "
-        "| ...........^",
-        str(e.exception),
-    )
+    if self.options.enable_pratt_parser:
+      self.assertIn(
+          "1:10: unexpected character\n"
+          " | 'Hello,' # 'World!'\n"
+          " | .........^",
+          str(e.exception),
+      )
+    else:
+      self.assertIn(
+          "1:10: Syntax error: token recognition error at: '#'\n "
+          "| 'Hello,' # 'World!'\n "
+          "| .........^",
+          str(e.exception),
+      )
+      self.assertIn(
+          "1:12: Syntax error: extraneous input ''World!'' expecting <EOF>\n "
+          "| 'Hello,' # 'World!'\n "
+          "| ...........^",
+          str(e.exception),
+      )
 
     # Check type-checker error.
     with self.assertRaises(Exception) as e:
@@ -793,7 +807,11 @@ def testCompilationErrorHandling(self):
     )
 
   def testErrorHandling(self):
-    bad_env = cel.NewEnv(_BadDescriptorPool(), variables={})
+    bad_env = cel.NewEnv(
+        _BadDescriptorPool(),
+        variables={},
+        options=self.options,
+    )
     with self.assertRaises(Exception) as e:
       bad_env.compile("cel.expr.conformance.proto2.TestSomeTypes{}")
     self.assertRegex(
@@ -929,5 +947,14 @@ def testErrorOnProtoCreation(self):
     )
 
 
+class CelTest(_CelTestBase):
+  # Default options.
+  pass
+
+
+class CelPrattParserTest(_CelTestBase):
+  options = cel.Options(enable_pratt_parser=True)
+
+
 if __name__ == "__main__":
   absltest.main()

cel_expr_python/py_cel_env.cc

@@ -32,6 +32,7 @@
 #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_options.h"
 #include "cel_expr_python/py_cel_type.h"
 #include "cel_expr_python/py_error_status.h"
 #include <pybind11/pybind11.h>
@@ -81,7 +82,8 @@ void PyCelEnv::DefinePythonBindings(pybind11::module& m) {
          std::optional<std::vector<std::shared_ptr<PyCelFunctionDecl>>>&
              functions,
          std::optional<std::unordered_map<std::string, py::object>>&
-             function_impls) {
+             function_impls,
+         std::optional<PyCelOptions>& options) {
         PyObject* pool_ptr;
         if (descriptor_pool.is_none()) {
           // Replicates python's `descriptor_pool.Default()`
@@ -119,7 +121,10 @@ void PyCelEnv::DefinePythonBindings(pybind11::module& m) {
           }
         }
 
-        return PyCelEnv(config.value_or(PyCelEnvConfig()), pool_ptr,
+        PyCelOptions env_options = options.value_or(PyCelOptions());
+
+        return PyCelEnv(config.value_or(PyCelEnvConfig()), env_options,
+                        pool_ptr,
                         std::move(variables).value_or(
                             std::unordered_map<std::string, PyCelType>{}),
                         ext_ptrs, std::move(expr_container),
@@ -131,10 +136,12 @@ void PyCelEnv::DefinePythonBindings(pybind11::module& m) {
       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("functions") = py::none(),
-      py::arg("function_impls") = py::none());
+      py::arg("function_impls") = py::none(), py::arg("options") = py::none());
   cel_class
       .def("config",
            [](PyCelEnv& self) { return self.GetEnv()->GetEnvConfig(); })
+      .def("options",
+           [](PyCelEnv& self) { return self.GetEnv()->GetOptions(); })
       .def("compile", &PyCelEnv::Compile, py::arg("expression"),
            py::arg("disable_check") = false)
       .def("deserialize", &PyCelEnv::Deserialize, py::arg("serialized"))
@@ -165,14 +172,15 @@ void PyCelEnv::DefinePythonBindings(pybind11::module& m) {
 }
 
 PyCelEnv::PyCelEnv(
-    const PyCelEnvConfig& config, PyObject* descriptor_pool,
+    const PyCelEnvConfig& config, const PyCelOptions& options,
+    PyObject* descriptor_pool,
     const std::unordered_map<std::string, PyCelType>& variable_types,
     const std::vector<PyObject*>& extensions,
     cel::ExpressionContainer 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,
+      config, options, descriptor_pool, std::move(variable_types), extensions,
       std::move(container), std::move(functions), std::move(function_impls)));
   ABSL_CHECK(PyGILState_Check());
 }

cel_expr_python/py_cel_env.h

@@ -30,6 +30,7 @@
 #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_options.h"
 #include "cel_expr_python/py_cel_type.h"
 #include <pybind11/pybind11.h>
 
@@ -68,7 +69,8 @@ class PyCelEnv {
 
  private:
   // Private constructor. Use `py_cel.NewEnv()` in python to obtain an instance.
-  PyCelEnv(const PyCelEnvConfig& config, PyObject* descriptor_pool,
+  PyCelEnv(const PyCelEnvConfig& config, const PyCelOptions& options,
+           PyObject* descriptor_pool,
            const std::unordered_map<std::string, PyCelType>& variable_types,
            const std::vector<PyObject*>& extensions,
            cel::ExpressionContainer container,

cel_expr_python/py_cel_env_internal.cc

@@ -45,6 +45,7 @@
 #include "cel_expr_python/py_cel_env_config.h"
 #include "cel_expr_python/py_cel_function.h"
 #include "cel_expr_python/py_cel_function_decl.h"
+#include "cel_expr_python/py_cel_options.h"
 #include "cel_expr_python/py_cel_overload.h"
 #include "cel_expr_python/py_cel_python_extension.h"
 #include "cel_expr_python/py_cel_type.h"
@@ -66,10 +67,12 @@ static const cel::FunctionDescriptorOptions kFunctionDescriptorOptions = {
 }  // namespace
 
 PyCelEnvInternal::PyCelEnvInternal(
-    const PyCelEnvConfig& env_config, PyObject* py_descriptor_pool,
+    const PyCelEnvConfig& env_config, const PyCelOptions& options,
+    PyObject* py_descriptor_pool,
     std::vector<CelExtensionHandle> extension_handles,
     absl::flat_hash_map<std::string, py::object>& function_impls)
     : env_config_(env_config),
+      options_(options),
       py_descriptor_database_(py_descriptor_pool),
       descriptor_pool_(
           std::make_shared<google::protobuf::DescriptorPool>(&py_descriptor_database_)),
@@ -105,7 +108,8 @@ PyCelEnvInternal::PyCelEnvInternal(
 
 absl::StatusOr<std::shared_ptr<PyCelEnvInternal>>
 PyCelEnvInternal::NewCelEnvInternal(
-    const PyCelEnvConfig& env_config, PyObject* py_descriptor_pool,
+    const PyCelEnvConfig& env_config, const PyCelOptions& options,
+    PyObject* py_descriptor_pool,
     const std::unordered_map<std::string, PyCelType>& variable_types,
     const std::vector<PyObject*>& extensions,
     cel::ExpressionContainer container,
@@ -219,7 +223,7 @@ PyCelEnvInternal::NewCelEnvInternal(
     }
   }
   return std::shared_ptr<PyCelEnvInternal>(
-      new PyCelEnvInternal(PyCelEnvConfig(config), py_descriptor_pool,
+      new PyCelEnvInternal(PyCelEnvConfig(config), options, py_descriptor_pool,
                            std::move(extension_handles), impls));
 }
 
@@ -237,6 +241,9 @@ absl::StatusOr<const cel::Compiler*> PyCelEnvInternal::GetCompiler(
       std::unique_ptr<cel::CompilerBuilder> compiler_builder,
       env->cel_env_.NewCompilerBuilder());
 
+  compiler_builder->GetParserBuilder().GetOptions().enable_pratt_parser =
+      env->options_.enable_pratt_parser;
+
   cel::TypeCheckerBuilder& checker_builder =
       compiler_builder->GetCheckerBuilder();
 

cel_expr_python/py_cel_env_internal.h

@@ -36,6 +36,7 @@
 #include "cel_expr_python/py_cel_env_config.h"
 #include "cel_expr_python/py_cel_function.h"
 #include "cel_expr_python/py_cel_function_decl.h"
+#include "cel_expr_python/py_cel_options.h"
 #include "cel_expr_python/py_cel_type.h"
 #include "cel_expr_python/py_descriptor_database.h"
 #include "cel_expr_python/py_message_factory.h"
@@ -73,14 +74,16 @@ class PyCelEnvInternal {
  public:
   ~PyCelEnvInternal() = default;
   static absl::StatusOr<std::shared_ptr<PyCelEnvInternal>> NewCelEnvInternal(
-      const PyCelEnvConfig& env_config, PyObject* py_descriptor_pool,
+      const PyCelEnvConfig& env_config, const PyCelOptions& options,
+      PyObject* py_descriptor_pool,
       const std::unordered_map<std::string, PyCelType>& variable_types,
       const std::vector<PyObject*>& extensions,
       cel::ExpressionContainer container,
       const std::vector<std::shared_ptr<PyCelFunctionDecl>>& functions,
       const std::unordered_map<std::string, py::object>& function_impls);
 
   const PyCelEnvConfig& GetEnvConfig() const { return env_config_; }
+  const PyCelOptions& GetOptions() const { return options_; }
 
   static absl::StatusOr<const cel::Compiler*> GetCompiler(
       const std::shared_ptr<PyCelEnvInternal>& env);
@@ -113,8 +116,8 @@ class PyCelEnvInternal {
  private:
   // Use NewCelEnvInternal() to create an instance.
   PyCelEnvInternal(
-      const PyCelEnvConfig& env_config, PyObject* py_descriptor_pool,
-      std::vector<CelExtensionHandle> extensions,
+      const PyCelEnvConfig& env_config, const PyCelOptions& options,
+      PyObject* py_descriptor_pool, std::vector<CelExtensionHandle> extensions,
       absl::flat_hash_map<std::string, py::object>& function_impls);
 
   absl::Status ConfigureStandardExtension(
@@ -128,6 +131,7 @@ class PyCelEnvInternal {
   cel::Env cel_env_;
   cel::EnvRuntime cel_env_runtime_;
   PyCelEnvConfig env_config_;
+  PyCelOptions options_;
   PyDescriptorDatabase py_descriptor_database_;
   std::shared_ptr<google::protobuf::DescriptorPool> descriptor_pool_;
   google::protobuf::DynamicMessageFactory message_factory_;

cel_expr_python/py_cel_module.cc

@@ -19,6 +19,7 @@
 #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_options.h"
 #include "cel_expr_python/py_cel_overload.h"
 #include "cel_expr_python/py_cel_python_extension.h"
 #include "cel_expr_python/py_cel_type.h"
@@ -39,6 +40,7 @@ PYBIND11_MODULE(cel, m) {
   PyCelFunctionDecl::DefinePythonBindings(m);
   PyCelPythonExtension::DefinePythonBindings(m);
   PyCelFunction::DefinePythonBindings(m);
+  PyCelOptions::DefinePythonBindings(m);
   PyCelEnvConfig::DefinePythonBindings(m);
   PyCelEnv::DefinePythonBindings(m);
 }