Improves handling on using methods as the values for callable symbolic attributes.

- Frozen fields will be absent from __init__ signature.
- Frozen fields will not be included during `format`.

PiperOrigin-RevId: 639901050

Author: daiyip

pyglove/core/symbolic/dict.py

@@ -796,6 +796,7 @@ def update(self,
 
   def sym_jsonify(
       self,
+      hide_frozen: bool = True,
       hide_default_values: bool = False,
       exclude_keys: Optional[Sequence[str]] = None,
       use_inferred: bool = False,
@@ -809,26 +810,30 @@ def sym_jsonify(
         # NOTE(daiyip): The key values of frozen field can safely be excluded
         # since they will be the same for a class.
         field = self._value_spec.schema[key_spec]
-        if not field.frozen:
-          for key in keys:
-            if key not in exclude_keys:
-              value = self.sym_getattr(key)
-              if use_inferred and isinstance(value, base.Inferential):
-                value = self.sym_inferred(key, default=value)
-              if pg_typing.MISSING_VALUE == value:
-                continue
-              if hide_default_values and base.eq(value, field.default_value):
-                continue
-              json_repr[key] = base.to_json(
-                  value, hide_default_values=hide_default_values,
-                  use_inferred=use_inferred,
-                  **kwargs)
+        if hide_frozen and field.frozen:
+          continue
+        for key in keys:
+          if key not in exclude_keys:
+            value = self.sym_getattr(key)
+            if use_inferred and isinstance(value, base.Inferential):
+              value = self.sym_inferred(key, default=value)
+            if pg_typing.MISSING_VALUE == value:
+              continue
+            if hide_default_values and base.eq(value, field.default_value):
+              continue
+            json_repr[key] = base.to_json(
+                value,
+                hide_frozen=hide_frozen,
+                hide_default_values=hide_default_values,
+                use_inferred=use_inferred,
+                **kwargs)
       return json_repr
     else:
       return {
           k: base.to_json(
               self.sym_inferred(k, default=v) if (
                   use_inferred and isinstance(v, base.Inferential)) else v,
+              hide_frozen=hide_frozen,
               hide_default_values=hide_default_values,
               use_inferred=use_inferred,
               **kwargs)
@@ -880,6 +885,7 @@ def format(
       *,
       python_format: bool = False,
       markdown: bool = False,
+      hide_frozen: bool = True,
       hide_default_values: bool = False,
       hide_missing_values: bool = False,
       include_keys: Optional[Set[str]] = None,
@@ -910,6 +916,8 @@ def _should_include_key(key):
         for key in keys:
           if _should_include_key(key):
             field = self._value_spec.schema[key_spec]
+            if hide_frozen and field.frozen:
+              continue
             v = self.sym_getattr(key)
             if use_inferred and isinstance(v, base.Inferential):
               v = self.sym_inferred(key, default=v)
@@ -941,6 +949,7 @@ def _should_include_key(key):
             compact,
             verbose,
             root_indent + 1,
+            hide_frozen=hide_frozen,
             hide_default_values=hide_default_values,
             hide_missing_values=hide_missing_values,
             python_format=python_format,
@@ -971,6 +980,7 @@ def _should_include_key(key):
             compact,
             verbose,
             root_indent + 1,
+            hide_frozen=hide_frozen,
             hide_default_values=hide_default_values,
             hide_missing_values=hide_missing_values,
             python_format=python_format,

pyglove/core/symbolic/dict_test.py

@@ -1907,6 +1907,36 @@ def __eq__(self, other):
     self.assertEqual(sd.to_json_str(), '{"x": 1, "y": 2.0}')
     self.assertEqual(base.from_json_str(sd.to_json_str(), value_spec=spec), sd)
 
+  def test_hide_frozen(self):
+
+    class A(pg_object.Object):
+      x: pg_typing.Int().freeze(1)
+
+    sd = Dict.partial(
+        a=A(),
+        value_spec=pg_typing.Dict([
+            ('a', pg_typing.Object(A)),
+            ('b', pg_typing.Bool(True).freeze()),
+        ]))
+    self.assertEqual(
+        sd.to_json(),
+        {
+            'a': {
+                '_type': A.__type_name__
+            },
+        }
+    )
+    self.assertEqual(
+        sd.to_json(hide_frozen=False),
+        {
+            'a': {
+                '_type': A.__type_name__,
+                'x': 1,
+            },
+            'b': True
+        }
+    )
+
   def test_hide_default_values(self):
 
     class A(pg_object.Object):
@@ -2255,6 +2285,24 @@ def test_noncompact_with_inferred_value(self):
         """),
     )
 
+  def test_hide_frozen(self):
+    d = Dict(x=1, value_spec=pg_typing.Dict([
+        ('x', pg_typing.Int()),
+        ('y', pg_typing.Bool(True).freeze()),
+        ('z', pg_typing.Dict([
+            ('v', pg_typing.Int(1)),
+            ('w', pg_typing.Bool().freeze(True)),
+        ]))
+    ]))
+    self.assertEqual(
+        d.format(compact=True),
+        '{x=1, z={v=1}}'
+    )
+    self.assertEqual(
+        d.format(compact=False, hide_frozen=False),
+        '{\n  x = 1,\n  y = True,\n  z = {\n    v = 1,\n    w = True\n  }\n}'
+    )
+
 
 def _on_change_callback(updates):
   del updates

pyglove/core/symbolic/object.py

@@ -149,22 +149,27 @@ def _infer_fields_from_annotations(cls) -> List[pg_typing.Field]:
     fields = cls._end_annotation_inference(fields)  # pytype: disable=attribute-error
     return fields
 
-  def _update_default_values_from_class_attributes(cls):
-    """Updates the symbolic attribute defaults from class attributes."""
-    for field in cls.__schema__.fields.values():
+  def _update_default_values_from_class_attributes(
+      cls, schema: pg_typing.Schema):
+    """Freezes callable fields if their defaults are provided as methods."""
+    for field in schema.fields.values():
       if isinstance(field.key, pg_typing.ConstStrKey):
-        attr_name = field.key.text
-        attr_value = cls.__dict__.get(attr_name, pg_typing.MISSING_VALUE)
-        if (
-            attr_value != pg_typing.MISSING_VALUE
-            and not isinstance(attr_value, property)
-            and (
-                # This allows class methods to be used as callable
-                # symbolic attributes.
-                not inspect.isfunction(attr_value)
-                or isinstance(field.value, pg_typing.Callable)
-                )
-        ):
+        attr_value = cls.__dict__.get(field.key.text, pg_typing.MISSING_VALUE)
+        if (attr_value == pg_typing.MISSING_VALUE
+            or isinstance(attr_value, property)):
+          continue
+        if inspect.isfunction(attr_value):
+          # When users add a method that has the same name as as field, two
+          # scenarios emerge. If the field is a callable type, the method will
+          # serve as the default value for the field. As a result, we freeze the
+          # field so it can't be provided from the constructor. If the field is
+          # not a callable type, the symbolic field and the method will coexist,
+          # meaning that the method has higher priority when being accessed,
+          # while users still can use `sym_getattr` to access the value for the
+          # symboic field.
+          if isinstance(field.value, pg_typing.Callable):
+            field.value.freeze(attr_value, apply_before_use=False)
+        else:
           field.value.set_default(attr_value)
 
 
@@ -320,6 +325,9 @@ def __init_subclass__(cls):
               metadata={},
           )
       )
+      # Freeze callable symbolic attributes if they are provided as methods.
+      user_cls._update_default_values_from_class_attributes(cls_schema)
+
       # NOTE(daiyip): When new fields are added through class attributes.
       # We invalidate `init_arg_list` so PyGlove could recompute it based
       # on its schema during `apply_schema`. Otherwise, we inherit the
@@ -333,10 +341,6 @@ def __init_subclass__(cls):
   @classmethod
   def _on_schema_update(cls):
     """Customizable trait: handling schema change."""
-    # Update the default value for each field after schema is updated. This is
-    # because that users may change a field's default value via class attribute.
-    cls._update_default_values_from_class_attributes()  # pylint: disable=no-value-for-parameter
-
     # Update all schema-based signatures.
     cls._update_signatures_based_on_schema()
 

pyglove/core/symbolic/object_test.py

@@ -107,7 +107,7 @@ def test_init_arg_list(self):
     self.assertEqual(
         self._A.init_arg_list, ['x', 'y', 'z', 'p'])
     self.assertEqual(
-        self._B.init_arg_list, ['x', 'y', 'z', 'p', 'q'])
+        self._B.init_arg_list, ['x', 'y', 'z', 'q'])
     self.assertEqual(
         self._C.init_arg_list, ['x', 'y', 'z', '*args'])
 
@@ -370,6 +370,9 @@ def x(self, v):
       def y(self):
         return self.sym_init_args.y * 2
 
+    self.assertTrue(H.__schema__.fields['x'].frozen)
+    self.assertFalse(H.__schema__.fields['y'].frozen)
+
     h = H(y=1)
     self.assertEqual(h.x(1), 3)
     self.assertEqual(h.y(), 2)

pyglove/core/symbolic/schema_utils.py

@@ -267,7 +267,8 @@ def auto_init_arg_list(cls):
   for base_cls in cls.__bases__:
     schema = getattr(base_cls, '__schema__', None)
     if isinstance(schema, pg_typing.Schema):
-      if list(schema.keys()) == list(cls.__schema__.keys()):
+      if ([(k, f.frozen) for k, f in schema.fields.items()]
+          == [(k, f.frozen) for k, f in cls.__schema__.fields.items()]):
         init_arg_list = base_cls.init_arg_list
       else:
         break
@@ -276,8 +277,8 @@ def auto_init_arg_list(cls):
     # declaration order from base classes to subclasses.
     init_arg_list = [
         str(key)
-        for key in cls.__schema__.fields.keys()
-        if isinstance(key, pg_typing.ConstStrKey)
+        for key, field in cls.__schema__.fields.items()
+        if isinstance(key, pg_typing.ConstStrKey) and not field.frozen
     ]
   return init_arg_list
 

pyglove/core/typing/callable_signature.py

@@ -228,7 +228,7 @@ def get_arg_spec(arg_name):
     kwonlyargs = []
     varkw = None
     for key, field in schema.fields.items():
-      if key not in existing_names:
+      if key not in existing_names and not field.frozen:
         if key.is_const:
           kwonlyargs.append(Argument(str(key), field.value))
         else:

pyglove/core/typing/callable_signature_test.py

@@ -230,6 +230,8 @@ def _get_signature(self, init_arg_list, is_method: bool = True):
     s = class_schema.Schema([
         class_schema.Field('x', vs.Int(), 'x'),
         class_schema.Field('y', vs.Int(), 'y'),
+        # Frozen fields will be ignored.
+        class_schema.Field('v', vs.Bool().freeze(True), 'v'),
         class_schema.Field('z', vs.List(vs.Int()), 'z'),
         class_schema.Field(ks.StrKey(), vs.Str(), 'kwargs'),
     ], metadata=dict(init_arg_list=init_arg_list), allow_nonconst_keys=True)