From ef6f82eaa24305b5ed9e97fe754867bd8719c4c7 Mon Sep 17 00:00:00 2001
From: Ryan McKenna <mckennar@google.com>
Date: Thu, 1 Oct 2026 13:08:51 -0700
Subject: [PATCH] Support optional `numerical_bounds` in
 `dpsynth.adapters.pydantic.infer_domain`.

Allows callers to supply or override `(min, max)` bounds by dot-separated field name (matching `dpsynth.adapters.protobuf.infer_domain`), which is useful when a shared nested `BaseModel` is reused across multiple fields with different domain ranges.

PiperOrigin-RevId: 991857818
---
 dpsynth/adapters/pydantic.py    | 21 ++++++++++++++++----
 tests/adapters/pydantic_test.py | 35 +++++++++++++++++++++++++++++++++
 2 files changed, 52 insertions(+), 4 deletions(-)

diff --git a/dpsynth/adapters/pydantic.py b/dpsynth/adapters/pydantic.py
index 0a245944..523cb0a6 100644
--- a/dpsynth/adapters/pydantic.py
+++ b/dpsynth/adapters/pydantic.py
@@ -68,6 +68,7 @@ def _get_base_type(annotation: type[Any]) -> tuple[bool, type[Any]]:
 
 def _numerical_attribute_from_field_info(
     field_info: FieldInfo,
+    bounds: tuple[float, float] | None = None,
 ) -> domain.NumericalAttribute:
   """Infers a NumericalAttribute from a pydantic FieldInfo."""
   # NumericalAttribute uses a convention where both the min_value and max_value
@@ -94,6 +95,8 @@ def _numerical_attribute_from_field_info(
           upper_bound = math.nextafter(meta.lt, -math.inf)
       case _:
         continue
+  if bounds is not None:
+    lower_bound, upper_bound = bounds
   if lower_bound is None or upper_bound is None:
     raise ValueError("Must specify lower and upper bounds for numeric fields.")
 
@@ -132,9 +135,11 @@ def _categorical_attribute_from_field_info(
 
 def infer_domain(
     model_cls: type[pydantic.BaseModel],
+    *,
+    numerical_bounds: Mapping[str, tuple[float, float]] | None = None,
 ) -> dict[str, domain.AttributeType]:
   """Infers the domain of a pydantic model."""
-
+  numerical_bounds = numerical_bounds or {}
   attributes: dict[str, domain.AttributeType] = {}
   for name, meta in model_cls.model_fields.items():
     _, base_type = _get_base_type(meta.annotation)  # pyrefly: ignore[bad-argument-type]
@@ -143,10 +148,18 @@ def infer_domain(
     is_model = is_class and issubclass(base_type, pydantic.BaseModel)
     is_literal = typing.get_origin(base_type) is Literal
     if is_model:
-      sub = infer_domain(base_type)
-      attributes.update({f"{name}.{k}": v for k, v in sub.items()})
+      prefix = f"{name}."
+      sub_bounds = {
+          k.removeprefix(prefix): v
+          for k, v in numerical_bounds.items()
+          if k.startswith(prefix)
+      }
+      sub = infer_domain(base_type, numerical_bounds=sub_bounds)
+      attributes.update({f"{prefix}{k}": v for k, v in sub.items()})
     elif base_type in (int, float):
-      attributes[name] = _numerical_attribute_from_field_info(meta)
+      attributes[name] = _numerical_attribute_from_field_info(
+          meta, bounds=numerical_bounds.get(name)
+      )
     elif base_type is str:
       attributes[name] = domain.OpenSetCategoricalAttribute(
           description=meta.description
diff --git a/tests/adapters/pydantic_test.py b/tests/adapters/pydantic_test.py
index 7a3b29b8..ae2038c9 100644
--- a/tests/adapters/pydantic_test.py
+++ b/tests/adapters/pydantic_test.py
@@ -238,6 +238,41 @@ def test_infer_domain_unsupported_type(self):
     ):
       pydantic_api.infer_domain(ModelWithUnsupportedType)
 
+  def test_infer_domain_with_numerical_bounds(self):
+    class Nested(pydantic.BaseModel):
+      val: float | None = pydantic.Field(description="Signal value")
+
+    class Root(pydantic.BaseModel):
+      steps: Nested
+      hrv: Nested
+
+    schema = pydantic_api.infer_domain(
+        Root,
+        numerical_bounds={
+            "steps.val": (0.0, 50000.0),
+            "hrv.val": (0.0, 300.0),
+        },
+    )
+    self.assertEqual(
+        schema,
+        {
+            "steps.val": domain.NumericalAttribute(
+                min_value=0.0,
+                max_value=50000.0,
+                clip_to_range=False,
+                dtype="float",
+                description="Signal value",
+            ),
+            "hrv.val": domain.NumericalAttribute(
+                min_value=0.0,
+                max_value=300.0,
+                clip_to_range=False,
+                dtype="float",
+                description="Signal value",
+            ),
+        },
+    )
+
   def test_open_set_str_description_and_optional_roundtrip(self):
     class ModelWithDescriptions(pydantic.BaseModel):
       timezone: str | None = pydantic.Field(description="IANA timezone")