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

pyglove/core/typing/class_schema.py

@@ -18,12 +18,12 @@ import copy
 import inspect
 import sys
 import types
-from typing import Any, Callable, Dict, Iterable, List, Optional, Set, Tuple, Type, Union
+from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Set, Tuple, Type, Union
 
-from pyglove.core import object_utils
+from pyglove.core import utils
 
 
-class KeySpec(object_utils.Formattable, object_utils.JSONConvertible):
+class KeySpec(utils.Formattable, utils.JSONConvertible):
   """Interface for key specifications.
 
   A key specification determines what keys are acceptable for a symbolic
@@ -94,7 +94,7 @@ class KeySpec(object_utils.Formattable, object_utils.JSONConvertible):
     assert False, 'Overridden in `key_specs.py`.'
 
 
-class ForwardRef(object_utils.Formattable):
+class ForwardRef(utils.Formattable):
   """Forward type reference."""
 
   def __init__(self, module: types.ModuleType, name: str):
@@ -139,14 +139,24 @@ class ForwardRef(object_utils.Formattable):
       )
     return reference
 
-  def format(self, *args, markdown: bool = False, **kwargs) -> str:
+  def format(
+      self,
+      compact: bool = False,
+      verbose: bool = True,
+      root_indent: int = 0,
+      **kwargs
+  ) -> str:
     """Format this object."""
-    details = object_utils.kvlist_str([
-        ('module', self.module.__name__, None),
-        ('name', self.name, None),
-    ])
-    return object_utils.maybe_markdown_quote(
-        f'{self.__class__.__name__}({details})', markdown
+    return utils.kvlist_str(
+        [
+            ('module', self.module.__name__, None),
+            ('name', self.name, None),
+        ],
+        label=self.__class__.__name__,
+        compact=compact,
+        verbose=verbose,
+        root_indent=root_indent,
+        **kwargs,
     )
 
   def __eq__(self, other: Any) -> bool:
@@ -156,7 +166,7 @@ class ForwardRef(object_utils.Formattable):
     elif isinstance(other, ForwardRef):
       return self.module is other.module and self.name == other.name
     elif inspect.isclass(other):
-      return self.resolved and self.cls is other
+      return self.resolved and self.cls is other  # pytype: disable=bad-return-type
 
   def __ne__(self, other: Any) -> bool:
     """Operator!=."""
@@ -170,7 +180,7 @@ class ForwardRef(object_utils.Formattable):
     return ForwardRef(self.module, self.name)
 
 
-class ValueSpec(object_utils.Formattable, object_utils.JSONConvertible):
+class ValueSpec(utils.Formattable, utils.JSONConvertible):
   """Interface for value specifications.
 
   A value specification defines what values are acceptable for a symbolic
@@ -334,6 +344,10 @@ class ValueSpec(object_utils.Formattable, object_utils.JSONConvertible):
       Tuple[Type[Any], ...]]:  # pyformat: disable
     """Returns acceptable (resolved) value type(s)."""
 
+  @abc.abstractmethod
+  def __call__(self, *args, **kwargs) -> Any:
+    """Instantiates a value based on the spec.."""
+
   @property
   @abc.abstractmethod
   def forward_refs(self) -> Set[ForwardRef]:
@@ -349,15 +363,19 @@ class ValueSpec(object_utils.Formattable, object_utils.JSONConvertible):
     """Returns True if current value spec accepts None."""
 
   @abc.abstractmethod
-  def set_default(self,
-                  default: Any,
-                  use_default_apply: bool = True) -> 'ValueSpec':
+  def set_default(
+      self,
+      default: Any,
+      use_default_apply: bool = True,
+      root_path: Optional[utils.KeyPath] = None,
+  ) -> 'ValueSpec':
     """Sets the default value and returns `self`.
 
     Args:
       default: Default value.
       use_default_apply: If True, invoke `apply` to the value, otherwise use
         default value as is.
+      root_path: (Optional) The path of the field.
 
     Returns:
       ValueSpec itself.
@@ -380,13 +398,14 @@ class ValueSpec(object_utils.Formattable, object_utils.JSONConvertible):
   @property
   def has_default(self) -> bool:
     """Returns True if the default value is provided."""
-    return self.default != object_utils.MISSING_VALUE
+    return self.default != utils.MISSING_VALUE
 
   @abc.abstractmethod
   def freeze(
       self,
-      permanent_value: Any = object_utils.MISSING_VALUE,
-      apply_before_use: bool = True) -> 'ValueSpec':
+      permanent_value: Any = utils.MISSING_VALUE,
+      apply_before_use: bool = True,
+  ) -> 'ValueSpec':
     """Sets the default value using a permanent value and freezes current spec.
 
     A frozen value spec will not accept any value that is not the default
@@ -453,10 +472,11 @@ class ValueSpec(object_utils.Formattable, object_utils.JSONConvertible):
       self,
       value: Any,
       allow_partial: bool = False,
-      child_transform: Optional[Callable[
-          [object_utils.KeyPath, 'Field', Any], Any]] = None,
-      root_path: Optional[object_utils.KeyPath] = None,
-      ) -> Any:
+      child_transform: Optional[
+          Callable[[utils.KeyPath, 'Field', Any], Any]
+      ] = None,
+      root_path: Optional[utils.KeyPath] = None,
+  ) -> Any:
     """Validates, completes and transforms the input value.
 
     Here is the procedure of ``apply``::
@@ -524,14 +544,16 @@ class ValueSpec(object_utils.Formattable, object_utils.JSONConvertible):
   def from_annotation(
       cls,
       annotation: Any,
-      auto_typing=False,
-      accept_value_as_annotation=False) -> 'ValueSpec':
+      auto_typing: bool = False,
+      accept_value_as_annotation: bool = False,
+      parent_module: Optional[types.ModuleType] = None
+      ) -> 'ValueSpec':
     """Gets a concrete ValueSpec from annotation."""
-    del annotation
+    del annotation, auto_typing, accept_value_as_annotation, parent_module
     assert False, 'Overridden in `annotation_conversion.py`.'
 
 
-class Field(object_utils.Formattable, object_utils.JSONConvertible):
+class Field(utils.Formattable, utils.JSONConvertible):
   """Class that represents the definition of one or a group of attributes.
 
   ``Field`` is held by a :class:`pyglove.Schema` object for defining the
@@ -661,9 +683,11 @@ class Field(object_utils.Formattable, object_utils.JSONConvertible):
       self,
       value: Any,
       allow_partial: bool = False,
-      transform_fn: Optional[Callable[
-          [object_utils.KeyPath, 'Field', Any], Any]] = None,
-      root_path: Optional[object_utils.KeyPath] = None) -> Any:
+      transform_fn: Optional[
+          Callable[[utils.KeyPath, 'Field', Any], Any]
+      ] = None,
+      root_path: Optional[utils.KeyPath] = None,
+  ) -> Any:
     """Apply current field to a value, which validate and complete the value.
 
     Args:
@@ -690,7 +714,8 @@ class Field(object_utils.Formattable, object_utils.JSONConvertible):
         value,
         allow_partial=allow_partial,
         child_transform=transform_fn,
-        root_path=root_path)
+        root_path=root_path
+    )
 
     if transform_fn:
       value = transform_fn(root_path, self, value)
@@ -711,34 +736,22 @@ class Field(object_utils.Formattable, object_utils.JSONConvertible):
       compact: bool = False,
       verbose: bool = True,
       root_indent: int = 0,
-      *,
-      markdown: bool = False,
       **kwargs,
   ) -> str:
     """Format this field into a string."""
-    description = self._description
-    if not verbose and self._description and len(self._description) > 20:
-      description = self._description[:20] + '...'
-
-    metadata = object_utils.format(
-        self._metadata,
+    return utils.kvlist_str(
+        [
+            ('key', self._key, None),
+            ('value', self._value, None),
+            ('description', self._description, None),
+            ('metadata', self._metadata, {}),
+        ],
+        label=self.__class__.__name__,
         compact=compact,
         verbose=verbose,
-        root_indent=root_indent + 1,
-        **kwargs)
-    if not verbose and len(metadata) > 24:
-      metadata = '{...}'
-    attr_str = object_utils.kvlist_str([
-        ('key', self._key, None),
-        ('value', self._value.format(
-            compact=compact,
-            verbose=verbose,
-            root_indent=root_indent + 1,
-            **kwargs), None),
-        ('description', object_utils.quote_if_str(description), None),
-        ('metadata', metadata, '{}')
-    ])
-    return object_utils.maybe_markdown_quote(f'Field({attr_str})', markdown)
+        root_indent=root_indent,
+        **kwargs,
+    )
 
   def to_json(self, **kwargs: Any) -> Dict[str, Any]:
     return self.to_json_dict(
@@ -766,7 +779,7 @@ class Field(object_utils.Formattable, object_utils.JSONConvertible):
     return not self.__eq__(other)
 
 
-class Schema(object_utils.Formattable, object_utils.JSONConvertible):
+class Schema(utils.Formattable, utils.JSONConvertible):
   """Class that represents a schema.
 
   PyGlove's runtime type system is based on the concept of ``Schema`` (
@@ -897,15 +910,51 @@ class Schema(object_utils.Formattable, object_utils.JSONConvertible):
         break
 
     if base_schema_list:
-      # Extend base schema from the nearest ancestor to the farthest.
-      for base in reversed(base_schema_list):
-        self.extend(base)
+      base = Schema.merge(base_schema_list)
+      self.extend(base)
 
     if not allow_nonconst_keys and self._dynamic_field is not None:
       raise ValueError(
           f'NonConstKey is not allowed in schema. '
           f'Encountered \'{self._dynamic_field.key}\'.')
 
+  @classmethod
+  def merge(
+      cls,
+      schema_list: Sequence['Schema'],
+      name: Optional[str] = None,
+      description: Optional[str] = None
+    ) -> 'Schema':
+    """Merge multiple schemas into one.
+
+    For fields shared by multiple schemas, the first appeared onces will be
+    used in the merged schema.
+
+    Args:
+      schema_list: A list of schemas to merge.
+      name: (Optional) name of the merged schema.
+      description: (Optinoal) description of the schema.
+
+    Returns:
+      The merged schema.
+    """
+    field_names = set()
+    fields = []
+    kw_field = None
+    for schema in schema_list:
+      for key, field in schema.fields.items():
+        if key.is_const and key not in field_names:
+          fields.append(field)
+          field_names.add(key)
+        elif not key.is_const and kw_field is None:
+          kw_field = field
+
+    if kw_field is not None:
+      fields.append(kw_field)
+    return Schema(
+        fields, name=name, description=description, allow_nonconst_keys=True
+    )
+
   def extend(self, base: 'Schema') -> 'Schema':
     """Extend current schema based on a base schema."""
 
@@ -914,13 +963,12 @@ class Schema(object_utils.Formattable, object_utils.JSONConvertible):
         parent_field: Field,
         child_field: Field) -> Field:
       """Merge function on field with the same key."""
-      if parent_field != object_utils.MISSING_VALUE:
-        if object_utils.MISSING_VALUE == child_field:
+      if parent_field != utils.MISSING_VALUE:
+        if utils.MISSING_VALUE == child_field:
           if (not self._allow_nonconst_keys and not parent_field.key.is_const):
-            hints = object_utils.kvlist_str([
-                ('base', object_utils.quote_if_str(base.name), None),
-                ('path', path, None)
-            ])
+            hints = utils.kvlist_str(
+                [('base', base.name, None), ('path', path, None)]
+            )
             raise ValueError(
                 f'Non-const key {parent_field.key} is not allowed to be '
                 f'added to the schema. ({hints})')
@@ -929,16 +977,15 @@ class Schema(object_utils.Formattable, object_utils.JSONConvertible):
           try:
             child_field.extend(parent_field)
           except Exception as e:  # pylint: disable=broad-except
-            hints = object_utils.kvlist_str([
-                ('base', object_utils.quote_if_str(base.name), None),
-                ('path', path, None)
-            ])
+            hints = utils.kvlist_str(
+                [('base', base.name, None), ('path', path, None)]
+            )
             raise e.__class__(f'{e} ({hints})').with_traceback(
                 sys.exc_info()[2])
       return child_field
 
-    self._fields = object_utils.merge([base.fields, self.fields], _merge_field)
-    self._metadata = object_utils.merge([base.metadata, self.metadata])
+    self._fields = utils.merge([base.fields, self.fields], _merge_field)
+    self._metadata = utils.merge([base.metadata, self.metadata])
 
     # Inherit dynamic field from base if it's not present in the child.
     if self._dynamic_field is None:
@@ -1061,8 +1108,8 @@ class Schema(object_utils.Formattable, object_utils.JSONConvertible):
       dict_obj: Dict[str, Any],
       allow_partial: bool = False,
       child_transform: Optional[Callable[
-          [object_utils.KeyPath, Field, Any], Any]] = None,
-      root_path: Optional[object_utils.KeyPath] = None,
+          [utils.KeyPath, Field, Any], Any]] = None,
+      root_path: Optional[utils.KeyPath] = None,
   ) -> Dict[str, Any]:  # pyformat: disable
     # pyformat: disable
     """Apply this schema to a dict object, validate and transform it.
@@ -1109,7 +1156,8 @@ class Schema(object_utils.Formattable, object_utils.JSONConvertible):
     if unmatched_keys:
       raise KeyError(
           f'Keys {unmatched_keys} are not allowed in Schema. '
-          f'(parent=\'{root_path}\')')
+          f'(parent=\'{root_path}\')'
+      )
 
     for key_spec, keys in matched_keys.items():
       field = self._fields[key_spec]
@@ -1118,19 +1166,19 @@ class Schema(object_utils.Formattable, object_utils.JSONConvertible):
         keys.append(str(key_spec))
       for key in keys:
         if dict_obj:
-          value = dict_obj.get(key, object_utils.MISSING_VALUE)
+          value = dict_obj.get(key, utils.MISSING_VALUE)
         else:
-          value = object_utils.MISSING_VALUE
+          value = utils.MISSING_VALUE
         # NOTE(daiyip): field.default_value may be MISSING_VALUE too
         # or partial.
-        if object_utils.MISSING_VALUE == value:
+        if utils.MISSING_VALUE == value:
           value = copy.deepcopy(field.default_value)
-
         new_value = field.apply(
             value,
             allow_partial=allow_partial,
             transform_fn=child_transform,
-            root_path=object_utils.KeyPath(key, root_path))
+            root_path=utils.KeyPath(key, root_path),
+        )
 
         # NOTE(daiyip): `pg.Dict.__getitem__`` has special logics in handling
         # `pg.Contextual`` values. Therefore, we user `dict.__getitem__()`` to
@@ -1143,10 +1191,12 @@ class Schema(object_utils.Formattable, object_utils.JSONConvertible):
           dict_obj[key] = new_value
     return dict_obj
 
-  def validate(self,
-               dict_obj: Dict[str, Any],
-               allow_partial: bool = False,
-               root_path: Optional[object_utils.KeyPath] = None) -> None:
+  def validate(
+      self,
+      dict_obj: Dict[str, Any],
+      allow_partial: bool = False,
+      root_path: Optional[utils.KeyPath] = None,
+  ) -> None:
     """Validates whether dict object is conformed with the schema."""
     self.apply(
         copy.deepcopy(dict_obj),
@@ -1210,50 +1260,27 @@ class Schema(object_utils.Formattable, object_utils.JSONConvertible):
       verbose: bool = True,
       root_indent: int = 0,
       *,
-      markdown: bool = False,
       cls_name: Optional[str] = None,
-      bracket_type: object_utils.BracketType = object_utils.BracketType.ROUND,
+      bracket_type: utils.BracketType = utils.BracketType.ROUND,
+      fields_only: bool = False,
       **kwargs,
   ) -> str:
     """Format current Schema into nicely printed string."""
-    if cls_name is None:
-      cls_name = 'Schema'
-
-    def _indent(text, indent):
-      return ' ' * 2 * indent + text
-
-    def _format_child(child):
-      return child.format(
-          compact=compact,
-          verbose=verbose,
-          root_indent=root_indent + 1,
-          **kwargs)
-
-    open_bracket, close_bracket = object_utils.bracket_chars(bracket_type)
-    if compact:
-      s = [f'{cls_name}{open_bracket}']
-      s.append(', '.join([
-          f'{f.key}={_format_child(f.value)}'
-          for f in self.fields.values()
-      ]))
-      s.append(close_bracket)
-    else:
-      s = [f'{cls_name}{open_bracket}\n']
-      last_field_show_description = False
-      for i, f in enumerate(self.fields.values()):
-        this_field_show_description = verbose and f.description
-        if i != 0:
-          s.append(',\n')
-          if last_field_show_description or this_field_show_description:
-            s.append('\n')
-        if this_field_show_description:
-          s.append(_indent(f'# {f.description}\n', root_indent + 1))
-        last_field_show_description = this_field_show_description
-        s.append(
-            _indent(f'{f.key} = {_format_child(f.value)}', root_indent + 1))
-      s.append('\n')
-      s.append(_indent(close_bracket, root_indent))
-    return object_utils.maybe_markdown_quote(''.join(s), markdown)
+    return utils.kvlist_str(
+        [
+            ('name', self.name, None),
+            ('description', self.description, None),
+            ('fields', list(self.fields.values()), []),
+            ('allow_nonconst_keys', self.allow_nonconst_keys, True),
+            ('metadata', self.metadata, {}),
+        ],
+        label=cls_name or self.__class__.__name__,
+        bracket_type=bracket_type,
+        compact=compact,
+        verbose=verbose,
+        root_indent=root_indent,
+        **kwargs,
+    )
 
   def to_json(self, **kwargs) -> Dict[str, Any]:
     return self.to_json_dict(
@@ -1277,15 +1304,41 @@ class Schema(object_utils.Formattable, object_utils.JSONConvertible):
     return not self.__eq__(other)
 
 
+FieldDef = Union[
+    # Key, Value spec/annotation.
+    Tuple[Union[str, KeySpec], Any],
+
+    # Key, Value spec/annotation, field docstr.
+    Tuple[Union[str, KeySpec], Any, str],
+
+    # Key, Value spec/annotation, field docstr, field metadata.
+    Tuple[Union[str, KeySpec], Any, str, Dict[str, Any]],
+]
+
+FieldKeyDef = Union[str, KeySpec]
+
+FieldValueDef = Union[
+    # Value spec/annotation.
+    Any,
+
+    # Value spec/annotation, field docstr.
+    Tuple[Any, str],
+
+    # Value spec/annotation, field docstr, field metadata.
+    Tuple[Any, str, Dict[str, Any]]
+]
+
+
 def create_field(
-    maybe_field: Union[Field, Tuple],   # pylint: disable=g-bare-generic
+    field_or_def: Union[Field, FieldDef],
     auto_typing: bool = True,
-    accept_value_as_annotation: bool = True
+    accept_value_as_annotation: bool = True,
+    parent_module: Optional[types.ModuleType] = None
 ) -> Field:
   """Creates ``Field`` from its equivalence.
 
   Args:
-    maybe_field: a ``Field`` object or its equivalence, which is a tuple of
+    field_or_def: a ``Field`` object or its equivalence, which is a tuple of
       2 - 4 elements:
       `(<key>, <value>, [description], [metadata])`.
       `key` can be a KeySpec subclass object or string. `value` can be a
@@ -1298,31 +1351,34 @@ def create_field(
       ``pg.typing.Any()`` will be used.
     accept_value_as_annotation: If True, allow default values to be used as
       annotations when creating the value spec.
+    parent_module: (Optional) parent module for defining this field, which will
+      be used for forward reference lookup.
 
   Returns:
     A ``Field`` object.
   """
-  if isinstance(maybe_field, Field):
-    return maybe_field
+  if isinstance(field_or_def, Field):
+    return field_or_def
 
-  if not isinstance(maybe_field, tuple):
+  if not isinstance(field_or_def, tuple):
     raise TypeError(
         f'Field definition should be tuples with 2 to 4 elements. '
-        f'Encountered: {maybe_field}.')
+        f'Encountered: {field_or_def}.')
 
-  if len(maybe_field) == 4:
-    maybe_key_spec, maybe_value_spec, description, field_metadata = maybe_field
-  elif len(maybe_field) == 3:
-    maybe_key_spec, maybe_value_spec, description = maybe_field
+  field_def = list(field_or_def)
+  if len(field_def) == 4:
+    maybe_key_spec, maybe_value_spec, description, field_metadata = field_def
+  elif len(field_def) == 3:
+    maybe_key_spec, maybe_value_spec, description = field_def
     field_metadata = {}
-  elif len(maybe_field) == 2:
-    maybe_key_spec, maybe_value_spec = maybe_field
+  elif len(field_def) == 2:
+    maybe_key_spec, maybe_value_spec = field_def
     description = None
     field_metadata = {}
   else:
     raise TypeError(
         f'Field definition should be tuples with 2 to 4 elements. '
-        f'Encountered: {maybe_field}.')
+        f'Encountered: {field_or_def}.')
 
   if isinstance(maybe_key_spec, (str, KeySpec)):
     key = maybe_key_spec
@@ -1330,10 +1386,14 @@ def create_field(
     raise TypeError(
         f'The 1st element of field definition should be of '
         f'<class \'str\'> or KeySpec. Encountered: {maybe_key_spec}.')
+
   value = ValueSpec.from_annotation(
       maybe_value_spec,
       auto_typing=auto_typing,
-      accept_value_as_annotation=accept_value_as_annotation)
+      accept_value_as_annotation=accept_value_as_annotation,
+      parent_module=parent_module,
+  )
+
   if (description is not None and
       not isinstance(description, str)):
     raise TypeError(f'Description (the 3rd element) of field definition '
@@ -1347,14 +1407,15 @@ def create_field(
 
 def create_schema(
     fields: Union[
-        Dict[str, Any],
-        List[Union[Field, Tuple]]   # pylint: disable=g-bare-generic
+        Dict[str, FieldValueDef],
+        List[Union[Field, FieldDef]]   # pylint: disable=g-bare-generic
     ],
     name: Optional[str] = None,
     base_schema_list: Optional[List[Schema]] = None,
     allow_nonconst_keys: bool = False,
     metadata: Optional[Dict[str, Any]] = None,
     description: Optional[str] = None,
+    parent_module: Optional[types.ModuleType] = None
 ) -> Schema:
   """Creates ``Schema`` from a list of ``Field``s or equivalences.
 
@@ -1395,6 +1456,8 @@ def create_schema(
     allow_nonconst_keys: Whether to allow non const keys in schema.
     metadata: Optional dict of user objects as schema-level metadata.
     description: Optional description of the schema.
+    parent_module: (Optional) parent module for defining this schema, which will
+      be used for forward reference lookup.
 
   Returns:
     Schema object.
@@ -1402,30 +1465,42 @@ def create_schema(
   Raises:
     TypeError: If input type is incorrect.
   """
-  if isinstance(fields, dict):
-    normalized_fields = []
-    for k, v in fields.items():
-      if not isinstance(v, (tuple, list)):
-        v = (v,)
-      normalized_fields.append(tuple([k] + list(v)))
-    fields = normalized_fields
-
-  if not isinstance(fields, list):
-    raise TypeError(
-        'Schema definition should be a dict of field names to their '
-        'definitions, a list of `pg.typing.Field` objects or a list of tuples '
-        'in format (key, value, description, metadata).')
-
+  fields = _normalize_field_defs(fields)
   metadata = metadata or {}
   if not isinstance(metadata, dict):
     raise TypeError(f'Metadata of schema should be a dict. '
                     f'Encountered: {metadata}.')
 
   return Schema(
-      fields=[create_field(maybe_field) for maybe_field in fields],
+      fields=[create_field(field_or_def, parent_module=parent_module)
+              for field_or_def in fields],
       name=name,
       base_schema_list=base_schema_list,
       allow_nonconst_keys=allow_nonconst_keys,
       metadata=metadata,
       description=description,
   )
+
+
+def _normalize_field_defs(
+    fields: Union[
+        Dict[str, FieldValueDef],
+        List[Union[Field, FieldDef]]   # pylint: disable=g-bare-generic
+    ]
+) -> List[Union[Field, FieldDef]]:
+  """Normalizes field definitions."""
+  if isinstance(fields, dict):
+    normalized_fields = []
+    for k, v in fields.items():
+      if not isinstance(v, (tuple, list)):
+        v = (v,)
+      normalized_fields.append(tuple([k] + list(v)))
+    return normalized_fields    # pytype: disable=bad-return-type
+  elif not isinstance(fields, list):
+    raise TypeError(
+        'Schema definition should be a dict of field names to their '
+        'definitions, a list of `pg.typing.Field` objects or a list of tuples '
+        'in format (key, value, description, metadata). '
+        f'Encountered: {fields}.'
+    )
+  return fields