pyglove 0.4.5.dev20240319__py3-none-any.whl → 0.4.5.dev202501132210__py3-none-any.whl

pyglove/core/symbolic/object.py

@@ -17,15 +17,14 @@ import abc
 import functools
 import inspect
 import typing
-from typing import Any, Dict, Iterator, List, Optional, Sequence, Tuple, Union
+from typing import Any, Dict, Iterator, List, Optional, Sequence, Union
 
-from pyglove.core import logging
-from pyglove.core import object_utils
+from pyglove.core import coding
 from pyglove.core import typing as pg_typing
+from pyglove.core import utils
 from pyglove.core.symbolic import base
 from pyglove.core.symbolic import dict as pg_dict
 from pyglove.core.symbolic import flags
-from pyglove.core.symbolic import schema_utils
 
 
 class ObjectMeta(abc.ABCMeta):
@@ -50,44 +49,56 @@ class ObjectMeta(abc.ABCMeta):
     """
     return f'{cls.__module__}.{cls.__qualname__}'
 
-  def __getattr__(cls, name):
-    # NOTE(daiyip): For backward compatibility, we allows these names to
-    # be used as aliases to the canonical names if users do not override them.
-    if name == 'schema':
-      logging.warning(
-          '`pg.Object.schema` is deprecated and will be removed in future. '
-          'Please use `__schema__` instead.')
-      return cls.__schema__
-    elif name == 'type_name':
-      logging.warning(
-          '`pg.Object.type_name` is deprecated and will be removed in future. '
-          'Please use `__type_name__` instead.')
-      return cls.__type_name__
-    raise AttributeError(name)
-
   @property
   def init_arg_list(cls) -> List[str]:
     """Gets __init__ positional argument list."""
     return typing.cast(List[str], cls.__schema__.metadata['init_arg_list'])
 
-  def apply_schema(cls, schema: pg_typing.Schema) -> None:
+  def apply_schema(cls, schema: Optional[pg_typing.Schema] = None) -> None:
     """Applies a schema to a symbolic class.
 
     Args:
       schema: The schema that will be applied to class. If `cls` was attached
-        with an existing schema. The old schema will be dropped.
+        with an existing schema. The old schema will be dropped. If None, the
+        cls will update its signature and getters according to the (maybe
+        updated) old schema.
     """
-    setattr(cls, '__schema__', schema)
-    setattr(cls, '__sym_fields', pg_typing.Dict(schema))
+    # Formalize schema first.
+    if schema is not None:
+      schema = cls._normalize_schema(schema)  # pytype: disable=attribute-error
+      setattr(cls, '__schema__', schema)
+      setattr(cls, '__sym_fields', pg_typing.Dict(schema))
 
-    # Update `init_arg_list`` based on the updated schema.
-    init_arg_list = schema.metadata.get('init_arg_list', None)
-    if init_arg_list is None:
-      init_arg_list = schema_utils.auto_init_arg_list(cls)
-      cls.__schema__.metadata['init_arg_list'] = init_arg_list
-    schema_utils.validate_init_arg_list(init_arg_list, cls.__schema__)
     cls._on_schema_update()  # pytype: disable=attribute-error
 
+  def update_schema(
+      cls,
+      fields: List[
+          Union[
+              pg_typing.Field,
+              List[pg_typing.FieldDef],
+              Dict[pg_typing.FieldKeyDef, pg_typing.FieldValueDef],
+          ]
+      ],
+      extend: bool = True,
+      *,
+      init_arg_list: Optional[Sequence[str]] = None,
+      metadata: Optional[Dict[str, Any]] = None,
+  ) -> None:
+    """Updates the schema of the class."""
+    metadata = metadata or {}
+    if init_arg_list is None:
+      init_arg_list = metadata.pop('init_arg_list', None)
+
+    metadata['init_arg_list'] = init_arg_list
+    schema = pg_typing.create_schema(
+        fields=fields,
+        base_schema_list=[cls.__schema__] if extend else [],
+        allow_nonconst_keys=True,
+        metadata=metadata,
+    )
+    cls.apply_schema(schema)
+
   def register_for_deserialization(
       cls,
       serialization_key: Optional[str] = None,
@@ -105,7 +116,7 @@ class ObjectMeta(abc.ABCMeta):
 
     # Register class with 'type' property.
     for key in serialization_keys:
-      object_utils.JSONConvertible.register(
+      utils.JSONConvertible.register(
           key, cls, flags.is_repeated_class_registration_allowed()
       )
 
@@ -138,33 +149,49 @@ class ObjectMeta(abc.ABCMeta):
       if key is None:
         continue
 
+      # Skip class-level attributes that are not symbolic fields.
+      if typing.get_origin(attr_annotation) is typing.ClassVar:
+        continue
+
       field = pg_typing.Field.from_annotation(key, attr_annotation)
       if isinstance(key, pg_typing.ConstStrKey):
         attr_value = cls.__dict__.get(attr_name, pg_typing.MISSING_VALUE)
         if attr_value != pg_typing.MISSING_VALUE:
           field.value.set_default(attr_value)
+
+      if (field.value.frozen and
+          field.value.default is
+          pg_typing.value_specs._FROZEN_VALUE_PLACEHOLDER):  # pylint: disable=protected-access
+        raise TypeError(
+            f'Field {field.key!r} is marked as final but has no default value.'
+        )
       fields.append(field)
 
     # Trigger event so subclass could modify the fields.
     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)
 
 
@@ -282,7 +309,7 @@ class Object(base.Symbolic, metaclass=ObjectMeta):
     Args:
       user_cls: The source class that calls this class method.
     """
-    object_utils.ensure_explicit_method_override(
+    utils.ensure_explicit_method_override(
         cls.__init__,
         (
             '`pg.Object.__init__` is a PyGlove managed method. For setting up '
@@ -290,7 +317,8 @@ class Object(base.Symbolic, metaclass=ObjectMeta):
             '`_on_init()`. If you do have a need to override `__init__` and '
             'know the implications, please decorate your overridden method '
             'with `@pg.explicit_method_override`.'
-        ))
+        ),
+    )
 
     # Set `__serialization_key__` before JSONConvertible.__init_subclass__
     # is called.
@@ -311,15 +339,16 @@ class Object(base.Symbolic, metaclass=ObjectMeta):
           base_schema_list.append(base_schema)
 
       new_fields = user_cls._infer_fields_from_annotations()
-      cls_schema = schema_utils.formalize_schema(
-          pg_typing.create_schema(
-              new_fields,
-              name=user_cls.__type_name__,
-              base_schema_list=base_schema_list,
-              allow_nonconst_keys=True,
-              metadata={},
-          )
+      cls_schema = pg_typing.create_schema(
+          new_fields,
+          base_schema_list=base_schema_list,
+          allow_nonconst_keys=True,
+          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
@@ -330,12 +359,99 @@ class Object(base.Symbolic, metaclass=ObjectMeta):
         cls_schema.metadata['init_arg_list'] = None
       user_cls.apply_schema(cls_schema)
 
+  @classmethod
+  def _normalize_schema(cls, schema: pg_typing.Schema) -> pg_typing.Schema:
+    """Normalizes the schema before applying it."""
+
+    schema.set_name(cls.__type_name__)
+    docstr = utils.docstr(cls)
+    if docstr:
+      schema.set_description(docstr.description)
+
+    def _formalize_field(path: utils.KeyPath, node: Any) -> bool:
+      """Formalize field."""
+      if isinstance(node, pg_typing.Field):
+        field = node
+        if (not flags.is_empty_field_description_allowed()
+            and not field.description):
+          raise ValueError(
+              f'Field description must not be empty (path={path}).')
+
+        field.value.set_default(
+            field.apply(
+                field.default_value,
+                allow_partial=True,
+                transform_fn=base.symbolic_transform_fn(allow_partial=True)),
+            use_default_apply=False)
+        if isinstance(field.value, pg_typing.Dict):
+          if field.value.schema is not None:
+            field.value.schema.set_name(f'{schema.name}.{path.path}')
+            utils.traverse(
+                field.value.schema.fields, _formalize_field, None, path
+            )
+        elif isinstance(field.value, pg_typing.List):
+          _formalize_field(utils.KeyPath(0, path), field.value.element)
+        elif isinstance(field.value, pg_typing.Tuple):
+          for i, elem in enumerate(field.value.elements):
+            _formalize_field(utils.KeyPath(i, path), elem)
+        elif isinstance(field.value, pg_typing.Union):
+          for i, c in enumerate(field.value.candidates):
+            _formalize_field(
+                utils.KeyPath(i, path),
+                pg_typing.Field(field.key, c, 'Union sub-type.'),
+            )
+      return True
+
+    utils.traverse(schema.fields, _formalize_field)
+    return schema
+
+  @classmethod
+  def _finalize_init_arg_list(cls) -> List[str]:
+    """Finalizes init_arg_list based on schema."""
+    # Update `init_arg_list`` based on the updated schema.
+    init_arg_list = cls.__schema__.metadata.get('init_arg_list', None)
+    if init_arg_list is None:
+      # Inherit from the first non-empty base if they have the same signature.
+      # This allows to bypass interface-only bases.
+      for base_cls in cls.__bases__:
+        schema = getattr(base_cls, '__schema__', None)  # pylint: disable=redefined-outer-name
+        if isinstance(schema, pg_typing.Schema):
+          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
+      if init_arg_list is None:
+        # Automatically generate from the field definitions in their
+        # declaration order from base classes to subclasses.
+        init_arg_list = [
+            str(key)
+            for key, field in cls.__schema__.fields.items()
+            if isinstance(key, pg_typing.ConstStrKey) and not field.frozen
+        ]
+      cls.__schema__.metadata['init_arg_list'] = init_arg_list
+    else:
+      for i, arg in enumerate(init_arg_list):
+        is_vararg = False
+        if i == len(init_arg_list) - 1 and arg.startswith('*'):
+          arg = arg[1:]
+          is_vararg = True
+        field = cls.__schema__.get_field(arg)
+        if field is None:
+          raise TypeError(
+              f'Argument {arg!r} from `init_arg_list` is not defined as a '
+              f'symbolic field. init_arg_list={init_arg_list!r}.')
+        if is_vararg and not isinstance(field.value, pg_typing.List):
+          raise TypeError(
+              f'Variable positional argument {arg!r} should be declared with '
+              f'`pg.typing.List(...)`. Encountered {field.value!r}.')
+    return init_arg_list
+
   @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
+    # Finalize init_arg_list baesd on schema.
+    cls._finalize_init_arg_list()
 
     # Update all schema-based signatures.
     cls._update_signatures_based_on_schema()
@@ -363,7 +479,7 @@ class Object(base.Symbolic, metaclass=ObjectMeta):
     # Create a new `__init__` that passes through all the arguments to
     # in `pg.Object.__init__`. This is needed for each class to use different
     # signature.
-    @object_utils.explicit_method_override
+    @utils.explicit_method_override
     @functools.wraps(pseudo_init)
     def _init(self, *args, **kwargs):
       # We pass through the arguments to `Object.__init__` instead of
@@ -392,7 +508,7 @@ class Object(base.Symbolic, metaclass=ObjectMeta):
   def _create_sym_attribute(cls, attr_name, field):
     """Customizable trait: template of single symbolic attribute."""
     return property(
-        object_utils.make_function(
+        coding.make_function(
             attr_name,
             ['self'],
             [f"return self.sym_inferred('{attr_name}')"],
@@ -426,7 +542,9 @@ class Object(base.Symbolic, metaclass=ObjectMeta):
       json_value: Any,
       *,
       allow_partial: bool = False,
-      root_path: Optional[object_utils.KeyPath] = None) -> 'Object':
+      root_path: Optional[utils.KeyPath] = None,
+      **kwargs,
+  ) -> 'Object':
     """Class method that load an symbolic Object from a JSON value.
 
     Example::
@@ -463,24 +581,26 @@ class Object(base.Symbolic, metaclass=ObjectMeta):
       json_value: Input JSON value, only JSON dict is acceptable.
       allow_partial: Whether to allow elements of the list to be partial.
       root_path: KeyPath of loaded object in its object tree.
+      **kwargs: Additional keyword arguments to pass through.
 
     Returns:
       A symbolic Object instance.
     """
     return cls(allow_partial=allow_partial, root_path=root_path, **{
-        k: base.from_json(v, allow_partial=allow_partial)
+        k: base.from_json(v, allow_partial=allow_partial, **kwargs)
         for k, v in json_value.items()
     })
 
-  @object_utils.explicit_method_override
+  @utils.explicit_method_override
   def __init__(
       self,
       *args,
       allow_partial: bool = False,
       sealed: Optional[bool] = None,
-      root_path: Optional[object_utils.KeyPath] = None,
+      root_path: Optional[utils.KeyPath] = None,
       explicit_init: bool = False,
-      **kwargs):
+      **kwargs,
+  ):
     """Create an Object instance.
 
     Args:
@@ -522,8 +642,8 @@ class Object(base.Symbolic, metaclass=ObjectMeta):
     # Fill field_args and init_args from **kwargs.
     _, unmatched_keys = self.__class__.__schema__.resolve(list(kwargs.keys()))
     if unmatched_keys:
-      arg_phrase = object_utils.auto_plural(len(unmatched_keys), 'argument')
-      keys_str = object_utils.comma_delimited_str(unmatched_keys)
+      arg_phrase = utils.auto_plural(len(unmatched_keys), 'argument')
+      keys_str = utils.comma_delimited_str(unmatched_keys)
       raise TypeError(
           f'{self.__class__.__name__}.__init__() got unexpected '
           f'keyword {arg_phrase}: {keys_str}')
@@ -543,8 +663,8 @@ class Object(base.Symbolic, metaclass=ObjectMeta):
         field_args[vararg_name] = list(args[num_named_args:])
         args = args[:num_named_args]
       elif len(args) > len(init_arg_names):
-        arg_phrase = object_utils.auto_plural(len(init_arg_names), 'argument')
-        was_phrase = object_utils.auto_plural(len(args), 'was', 'were')
+        arg_phrase = utils.auto_plural(len(init_arg_names), 'argument')
+        was_phrase = utils.auto_plural(len(args), 'was', 'were')
         raise TypeError(
             f'{self.__class__.__name__}.__init__() takes '
             f'{len(init_arg_names)} positional {arg_phrase} but {len(args)} '
@@ -556,7 +676,7 @@ class Object(base.Symbolic, metaclass=ObjectMeta):
 
     for k, v in kwargs.items():
       if k in field_args:
-        values_str = object_utils.comma_delimited_str([field_args[k], v])
+        values_str = utils.comma_delimited_str([field_args[k], v])
         raise TypeError(
             f'{self.__class__.__name__}.__init__() got multiple values for '
             f'argument \'{k}\': {values_str}.')
@@ -571,8 +691,8 @@ class Object(base.Symbolic, metaclass=ObjectMeta):
             and field.key not in field_args):
           missing_args.append(str(field.key))
       if missing_args:
-        arg_phrase = object_utils.auto_plural(len(missing_args), 'argument')
-        keys_str = object_utils.comma_delimited_str(missing_args)
+        arg_phrase = utils.auto_plural(len(missing_args), 'argument')
+        keys_str = utils.comma_delimited_str(missing_args)
         raise TypeError(
             f'{self.__class__.__name__}.__init__() missing {len(missing_args)} '
             f'required {arg_phrase}: {keys_str}.')
@@ -622,8 +742,7 @@ class Object(base.Symbolic, metaclass=ObjectMeta):
     and during __init__.
     """
 
-  def _on_change(self,
-                 field_updates: Dict[object_utils.KeyPath, base.FieldUpdate]):
+  def _on_change(self, field_updates: Dict[utils.KeyPath, base.FieldUpdate]):
     """Event that is triggered when field values in the subtree are updated.
 
     This event will be called
@@ -643,8 +762,7 @@ class Object(base.Symbolic, metaclass=ObjectMeta):
     del field_updates
     return self._on_bound()
 
-  def _on_path_change(
-      self, old_path: object_utils.KeyPath, new_path: object_utils.KeyPath):
+  def _on_path_change(self, old_path: utils.KeyPath, new_path: utils.KeyPath):
     """Event that is triggered after the symbolic path changes."""
     del old_path, new_path
 
@@ -723,8 +841,8 @@ class Object(base.Symbolic, metaclass=ObjectMeta):
     return self._sym_attributes.sym_getattr(key)
 
   def _sym_rebind(
-      self, path_value_pairs: Dict[object_utils.KeyPath, Any]
-      ) -> List[base.FieldUpdate]:
+      self, path_value_pairs: Dict[utils.KeyPath, Any]
+  ) -> List[base.FieldUpdate]:
     """Rebind current object using object-form members."""
     if base.treats_as_sealed(self):
       raise base.WritePermissionError(
@@ -763,9 +881,8 @@ class Object(base.Symbolic, metaclass=ObjectMeta):
     return self
 
   def _update_children_paths(
-      self,
-      old_path: object_utils.KeyPath,
-      new_path: object_utils.KeyPath) -> None:
+      self, old_path: utils.KeyPath, new_path: utils.KeyPath
+  ) -> None:
     """Update children paths according to root_path of current node."""
     self._sym_attributes.sym_setpath(new_path)
     self._on_path_change(old_path, new_path)
@@ -849,16 +966,15 @@ class Object(base.Symbolic, metaclass=ObjectMeta):
       return self.sym_hash()
     return super().__hash__()
 
-  def sym_jsonify(self, **kwargs) -> object_utils.JSONValueType:
+  def sym_jsonify(self, **kwargs) -> utils.JSONValueType:
     """Converts current object to a dict of plain Python objects."""
-    return object_utils.merge([
-        {
-            object_utils.JSONConvertible.TYPE_NAME_KEY: (
-                self.__class__.__serialization_key__
-            )
-        },
-        self._sym_attributes.to_json(**kwargs),
-    ])
+    json_dict = {
+        utils.JSONConvertible.TYPE_NAME_KEY: (
+            self.__class__.__serialization_key__
+        )
+    }
+    json_dict.update(self._sym_attributes.to_json(**kwargs))
+    return json_dict
 
   def format(self,
              compact: bool = False,
@@ -872,21 +988,24 @@ class Object(base.Symbolic, metaclass=ObjectMeta):
         root_indent,
         cls_name=self.__class__.__name__,
         key_as_attribute=True,
-        bracket_type=object_utils.BracketType.ROUND,
-        **kwargs)
+        bracket_type=utils.BracketType.ROUND,
+        **kwargs,
+    )
 
 
 base.Symbolic.ObjectType = Object
 
 
 def members(
-    fields: List[Union[
-        pg_typing.Field,
-        Tuple[Union[str, pg_typing.KeySpec], pg_typing.ValueSpec, str],
-        Tuple[Union[str, pg_typing.KeySpec], pg_typing.ValueSpec, str, Any]]],
+    fields: Union[
+        List[Union[pg_typing.Field, pg_typing.FieldDef]],
+        Dict[pg_typing.FieldKeyDef, pg_typing.FieldValueDef],
+    ],
     metadata: Optional[Dict[str, Any]] = None,
     init_arg_list: Optional[Sequence[str]] = None,
-    **kwargs
+    serialization_key: Optional[str] = None,
+    additional_keys: Optional[List[str]] = None,
+    add_to_registry: bool = True,
 ) -> pg_typing.Decorator:
   """Function/Decorator for declaring symbolic fields for ``pg.Object``.
 
@@ -946,16 +1065,17 @@ def members(
       provided, the `init_arg_list` will be automatically generated from
       symbolic attributes defined from ``pg.members`` in their declaration
       order, from the base classes to the subclass.
-    **kwargs: Keyword arguments for infrequently used options. Acceptable
-      keywords are:  * `serialization_key`: An optional string to be used as the
+    serialization_key: An optional string to be used as the
       serialization key for the class during `sym_jsonify`. If None,
       `cls.__type_name__` will be used. This is introduced for scenarios when we
       want to relocate a class, before the downstream can recognize the new
-      location, we need the class to serialize it using previous key. *
-      `additional_keys`: An optional list of strings as additional keys to
+      location, we need the class to serialize it using previous key.
+    additional_keys: An optional list of strings as additional keys to
       deserialize an object of the registered class. This can be useful when we
       need to relocate or rename the registered class while being able to load
       existing serialized JSON values.
+    add_to_registry: If True, register serialization keys and additional keys
+      with the class.
 
   Returns:
     a decorator function that register the class or function with schema
@@ -967,22 +1087,16 @@ def members(
     KeyError: If type has already been registered in the registry.
     ValueError: schema cannot be created from fields.
   """
-  serialization_key = kwargs.pop('serialization_key', None)
-  additional_keys = kwargs.pop('additional_keys', None)
-  if kwargs:
-    raise TypeError(f'Unsupported keyword arguments: {list(kwargs.keys())!r}.')
-
   def _decorator(cls):
     """Decorator function that registers schema with an Object class."""
-    schema_utils.update_schema(
-        cls,
+    cls.update_schema(
         fields,
         extend=True,
         init_arg_list=init_arg_list,
         metadata=metadata,
-        serialization_key=serialization_key,
-        additional_keys=additional_keys,
     )
+    if add_to_registry:
+      cls.register_for_deserialization(serialization_key, additional_keys)
     return cls
   return typing.cast(pg_typing.Decorator, _decorator)
 
@@ -1014,8 +1128,6 @@ def use_init_args(init_arg_list: Sequence[str]) -> pg_typing.Decorator:
     a decorator function that updates the `__init__` signature.
   """
   def _decorator(cls):
-    schema_utils.update_schema(
-        cls, [], extend=True, init_arg_list=init_arg_list
-    )
+    cls.update_schema([], extend=True, init_arg_list=init_arg_list)
     return cls
   return typing.cast(pg_typing.Decorator, _decorator)