[cel-expr-python] Replace leak census with regrtest-style check.

cel_expr_python/BUILD

@@ -1,5 +1,6 @@
 load("@pybind11_bazel//:build_defs.bzl", "pybind_extension", "pybind_library")
 load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load("@rules_python//python:py_library.bzl", "py_library")
 load("@rules_python//python:py_test.bzl", "py_test")
 
 package(default_visibility = ["//visibility:public"])
@@ -150,13 +151,21 @@ cc_library(
     ],
 )
 
+py_library(
+    name = "leak_check",
+    testonly = True,
+    srcs = ["leak_check.py"],
+    visibility = [":__subpackages__"],
+)
+
 py_test(
     name = "cel_test",
     srcs = ["cel_test.py"],
     data = [
         ":cel",
     ],
     deps = [
+        ":leak_check",
         "//testing:proto2_test_all_types_py_pb2",
         "@com_google_absl_py//absl/testing:absltest",
         "@com_google_protobuf//:protobuf",
@@ -174,6 +183,7 @@ py_test(
         ":cel",
     ],
     deps = [
+        ":leak_check",
         "//testing:proto2_test_all_types_py_pb2",
         "@com_google_absl_py//absl/testing:absltest",
     ] + select({

cel_expr_python/cel_parallel_test.py

@@ -25,6 +25,7 @@
 from absl.testing import absltest
 from cel.expr.conformance.proto2 import test_all_types_pb2 as test_all_types_pb
 from cel_expr_python import cel
+from cel_expr_python import leak_check
 
 
 @dataclasses.dataclass(frozen=True)
@@ -94,6 +95,7 @@ class _TestCase:
 ]
 
 
+@leak_check.leak_checked
 class CelParallelTest(absltest.TestCase):
 
   def setUp(self):
@@ -119,7 +121,6 @@ def setUp(self):
             )
         ],
     )
-    self.object_counts_before_test = self._grab_object_counts()
 
   def tearDown(self):
     """Tears down the test environment."""
@@ -128,26 +129,6 @@ def tearDown(self):
     gc.collect()
     # Assert that all Arenas have been garbage-collected
     self.assertEqual(cel._InternalArena._get_instance_count(), 0)
-    self._check_for_leaks()
-
-  def _grab_object_counts(self) -> dict[str, int]:
-    gc.collect()
-    all_objects = gc.get_objects()
-    type_counts = {}
-    for obj in all_objects:
-      obj_type = type(obj)
-      type_counts[obj_type.__name__] = type_counts.get(obj_type, 0) + 1
-    return type_counts
-
-  def _check_for_leaks(self):
-    type_counts = self._grab_object_counts()
-    for key, count in type_counts.items():
-      if count != self.object_counts_before_test.get(key, 0):
-        self.fail(
-            f"Object count for {key} did not match expected count. "
-            f"Expected: {self.object_counts_before_test.get(key, 0)}, "
-            f"Actual: {count}",
-        )
 
   def _test_eval(self, multi_threaded: bool):
     compiled_exprs = [self.env.compile(tc.expr) for tc in _TEST_CASES]

cel_expr_python/cel_test.py

@@ -28,8 +28,10 @@
 from absl.testing import absltest
 from cel.expr.conformance.proto2 import test_all_types_pb2 as test_all_types_pb
 from cel_expr_python import cel
+from cel_expr_python import leak_check
 
 
+@leak_check.leak_checked
 @absltest.skipThisClass("Base class")
 class _CelTestBase(absltest.TestCase):
   options: cel.Options = cel.Options()
@@ -57,7 +59,6 @@ def setUp(self):
         },
         options=self.options,
     )
-    self.object_counts_before_test = self._grab_object_counts()
 
   def tearDown(self):
     """Tears down the test environment."""
@@ -66,26 +67,6 @@ def tearDown(self):
     gc.collect()
     # Assert that all Arenas have been garbage-collected
     self.assertEqual(cel._InternalArena._get_instance_count(), 0)
-    self._check_for_leaks()
-
-  def _grab_object_counts(self) -> dict[str, int]:
-    gc.collect()
-    all_objects = gc.get_objects()
-    type_counts = {}
-    for obj in all_objects:
-      obj_type = type(obj)
-      type_counts[obj_type.__name__] = type_counts.get(obj_type, 0) + 1
-    return type_counts
-
-  def _check_for_leaks(self):
-    type_counts = self._grab_object_counts()
-    for key, count in type_counts.items():
-      if count != self.object_counts_before_test.get(key, 0):
-        self.fail(
-            f"Object count for {key} did not match expected count. "
-            f"Expected: {self.object_counts_before_test.get(key, 0)}, "
-            f"Actual: {count}",
-        )
 
   def _eval(
       self,

cel_expr_python/leak_check.py

@@ -0,0 +1,137 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Reference and memory leak detection helpers for CEL Python tests."""
+
+import array
+import collections.abc
+import functools
+import gc
+import sys
+import tracemalloc
+from typing import Any, TypeVar
+
+_DEFAULT_WARMUPS = 2
+_DEFAULT_RUNS = 3
+
+_F = TypeVar("_F", bound=collections.abc.Callable[..., Any])
+_T = TypeVar("_T", bound=type)
+
+
+def _live_blocks() -> int:
+  """Returns the allocated Python memory block count after clearing caches."""
+  # Clear CPython's internal caches (e.g., the type attribute lookup cache)
+  # before sampling so one-time cached entries are not mistaken for leaks.
+  clear_caches = getattr(sys, "_clear_internal_caches", None) or getattr(
+      sys, "_clear_type_cache", None
+  )
+  if clear_caches is not None:
+    clear_caches()
+  gc.collect()
+  return sys.getallocatedblocks()
+
+
+def _tracemalloc_report(
+    run_once: collections.abc.Callable[[], None],
+    iterations: int = _DEFAULT_RUNS,
+    limit: int = 5,
+) -> str:
+  """Re-runs `run_once` under tracemalloc and formats top growing sites."""
+  was_tracing = tracemalloc.is_tracing()
+  if not was_tracing:
+    tracemalloc.start(10)
+  try:
+    run_once()
+    _live_blocks()
+    before = tracemalloc.take_snapshot()
+    for _ in range(iterations):
+      run_once()
+    _live_blocks()
+    after = tracemalloc.take_snapshot()
+  finally:
+    if not was_tracing:
+      tracemalloc.stop()
+  ignore = [tracemalloc.Filter(False, "<frozen importlib._bootstrap*>")]
+  if tracemalloc.__file__:
+    ignore.append(tracemalloc.Filter(False, tracemalloc.__file__))
+  stats = after.filter_traces(ignore).compare_to(
+      before.filter_traces(ignore), "traceback"
+  )
+  lines = []
+  sites = 0
+  for stat in stats:
+    if stat.count_diff <= 0:
+      continue
+    lines.append(
+        f"  {stat.count_diff:+d} blocks ({stat.size_diff:+d} B) allocated at:"
+    )
+    lines.extend(
+        "    " + line for line in stat.traceback.format(most_recent_first=True)
+    )
+    sites += 1
+    if sites >= limit:
+      break
+  return "\n".join(lines) or "  (no growth attributed by tracemalloc)"
+
+
+def assert_no_leaks(
+    test_method: _F,
+    *,
+    warmups: int = _DEFAULT_WARMUPS,
+    runs: int = _DEFAULT_RUNS,
+) -> _F:
+  """Wraps a `unittest.TestCase` method to fail if live blocks grow every run.
+
+  Uses CPython's `regrtest -R` strategy: runs `warmups` unmeasured iterations
+  to populate lazy process-wide caches, then `runs` measured iterations,
+  failing only if `sys.getallocatedblocks()` increases on every measured run.
+
+  Args:
+    test_method: The test method to wrap.
+    warmups: Number of warm-up repetitions before measuring.
+    runs: Number of measured repetitions.
+
+  Returns:
+    The wrapped test method.
+  """
+
+  @functools.wraps(test_method)
+  def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
+    total_runs = warmups + runs
+    # Store samples in a C array so recording them keeps no Python ints alive.
+    blocks = array.array("q", [0]) * (total_runs + 1)
+    blocks[0] = _live_blocks()
+    for i in range(total_runs):
+      test_method(self, *args, **kwargs)
+      blocks[i + 1] = _live_blocks()
+    deltas = [blocks[i + 1] - blocks[i] for i in range(warmups, total_runs)]
+    if all(d > 0 for d in deltas):
+      self.fail(
+          f"Memory blocks grew on every one of {runs} repetitions"
+          f" (after {warmups} warm-up runs): deltas={deltas}."
+          " Top growing allocation sites:\n"
+          + _tracemalloc_report(
+              lambda: test_method(self, *args, **kwargs), iterations=runs
+          )
+      )
+
+  return wrapper  # type: ignore[return-value]
+
+
+def leak_checked(cls: _T) -> _T:
+  """Class decorator wrapping every `test*` method with `assert_no_leaks`."""
+  for name, attr in list(cls.__dict__.items()):
+    if name.startswith("test") and callable(attr):
+      setattr(cls, name, assert_no_leaks(attr))
+  return cls

cel_expr_python/py_cel_value.cc

@@ -55,6 +55,24 @@ namespace cel_python {
 
 namespace py = ::pybind11;
 
+namespace {
+
+std::string PyObjectRepr(PyObject* obj) {
+  if (obj == nullptr) {
+    return "<null>";
+  }
+  PyObject* repr = PyObject_Repr(obj);
+  const char* repr_cstr = repr != nullptr ? PyUnicode_AsUTF8(repr) : nullptr;
+  if (repr_cstr == nullptr) {
+    PyErr_Clear();
+  }
+  std::string repr_str = repr_cstr != nullptr ? repr_cstr : "<unknown>";
+  Py_XDECREF(repr);
+  return repr_str;
+}
+
+}  // namespace
+
 void PyCelValue::DefinePythonBindings(py::module& m) {
   py::class_<PyCelValue>(m, "Value")
       .def("type", &PyCelValue::Type)
@@ -767,9 +785,9 @@ absl::StatusOr<cel::Value> PyObjectToCelValue(
           time = py::cast<absl::Time>(py_object);
         } catch (const py::cast_error& e) {
           PyErr_Clear();
-          return absl::UnknownError(absl::StrFormat(
-              "Cannot convert %s to absl::Time: %s",
-              PyUnicode_AsUTF8(PyObject_Repr(py_object)), e.what()));
+          return absl::UnknownError(
+              absl::StrFormat("Cannot convert %s to absl::Time: %s",
+                              PyObjectRepr(py_object), e.what()));
         }
         absl::StatusOr<cel::TimestampValue> timestamp_value =
             cel::SafeTimestampValue(time);
@@ -808,7 +826,7 @@ absl::StatusOr<cel::Value> PyObjectToCelValue(
           PyErr_Clear();
           return absl::UnknownError(
               absl::StrFormat("Cannot convert %s to cel.Type.TIMESTAMP",
-                              PyUnicode_AsUTF8(PyObject_Repr(py_object))));
+                              PyObjectRepr(py_object)));
         }
         absl::Time time =
             absl::FromUnixSeconds(seconds) + absl::Nanoseconds(nanos);
@@ -830,9 +848,9 @@ absl::StatusOr<cel::Value> PyObjectToCelValue(
           duration = py::cast<absl::Duration>(py_object);
         } catch (const py::cast_error& e) {
           PyErr_Clear();
-          return absl::UnknownError(absl::StrFormat(
-              "Cannot convert %s to absl::Duration: %s",
-              PyUnicode_AsUTF8(PyObject_Repr(py_object)), e.what()));
+          return absl::UnknownError(
+              absl::StrFormat("Cannot convert %s to absl::Duration: %s",
+                              PyObjectRepr(py_object), e.what()));
         }
         absl::StatusOr<cel::DurationValue> duration_value =
             cel::SafeDurationValue(duration);
@@ -869,9 +887,8 @@ absl::StatusOr<cel::Value> PyObjectToCelValue(
         Py_XDECREF(py_nanos);
         if (!ok) {
           PyErr_Clear();
-          return absl::UnknownError(
-              absl::StrFormat("Cannot convert %s to absl::Duration",
-                              PyUnicode_AsUTF8(PyObject_Repr(py_object))));
+          return absl::UnknownError(absl::StrFormat(
+              "Cannot convert %s to absl::Duration", PyObjectRepr(py_object)));
         }
         absl::Duration duration =
             absl::Seconds(seconds) + absl::Nanoseconds(nanos);
@@ -953,28 +970,26 @@ absl::StatusOr<cel::Value> PyObjectToCelValue(
         while (PyDict_Next(py_object, &pos, &key, &value)) {
           // Note: PyDict_Next returns borrowed references, so we shouldn't
           // DECREF them.
-          CEL_PYTHON_ASSIGN_OR_RETURN(
-              cel::Value converted_key,
-              PyObjectToCelValue(
-                  key, key_type,
-                  [context, key]() {
-                    return absl::StrFormat(
-                        "%s[%s]", context(),
-                        PyUnicode_AsUTF8(PyObject_Repr(key)));
-                  },
-                  env, arena));
-
-          CEL_PYTHON_ASSIGN_OR_RETURN(
-              cel::Value converted_value,
-              PyObjectToCelValue(
-                  value, value_type,
-                  [context, key, value]() {
-                    return absl::StrFormat(
-                        "%s[%s] -> %s", context(),
-                        PyUnicode_AsUTF8(PyObject_Repr(key)),
-                        PyUnicode_AsUTF8(PyObject_Repr(value)));
-                  },
-                  env, arena));
+          CEL_PYTHON_ASSIGN_OR_RETURN(cel::Value converted_key,
+                                      PyObjectToCelValue(
+                                          key, key_type,
+                                          [context, key]() {
+                                            return absl::StrFormat(
+                                                "%s[%s]", context(),
+                                                PyObjectRepr(key));
+                                          },
+                                          env, arena));
+
+          CEL_PYTHON_ASSIGN_OR_RETURN(cel::Value converted_value,
+                                      PyObjectToCelValue(
+                                          value, value_type,
+                                          [context, key, value]() {
+                                            return absl::StrFormat(
+                                                "%s[%s] -> %s", context(),
+                                                PyObjectRepr(key),
+                                                PyObjectRepr(value));
+                                          },
+                                          env, arena));
           CEL_PYTHON_RETURN_IF_ERROR(
               builder->Put(converted_key, converted_value));
         }