pyglove 0.5.0.dev202510150810__py3-none-any.whl → 0.5.0.dev202510170226__py3-none-any.whl

pyglove/core/typing/class_schema.py

@@ -582,6 +582,16 @@ class ValueSpec(utils.Formattable, utils.JSONConvertible):
     del include_type_name, include_subclasses, inline_nested_refs, kwargs
     assert False, 'Overridden in `json_schema.py`.'
 
+  @classmethod
+  def from_json_schema(
+      cls,
+      json_schema: Dict[str, Any],
+      class_fn: Optional[Callable[[str, 'Schema'], Type[Any]]] = None
+  ) -> 'ValueSpec':
+    """Converts a JSON schema to a value spec."""
+    del json_schema, class_fn
+    assert False, 'Overridden in `json_schema.py`.'
+
 
 class Field(utils.Formattable, utils.JSONConvertible):
   """Class that represents the definition of one or a group of attributes.
@@ -1375,6 +1385,16 @@ class Schema(utils.Formattable, utils.JSONConvertible):
   def __ne__(self, other: Any) -> bool:
     return not self.__eq__(other)
 
+  @classmethod
+  def from_json_schema(
+      cls,
+      json_schema: Dict[str, Any],
+      class_fn: Optional[Callable[[str, 'Schema'], Type[Any]]] = None
+  ) -> 'Schema':
+    """Converts a JSON schema to a schema."""
+    del json_schema, class_fn
+    assert False, 'Overridden in `json_schema.py`.'
+
 
 FieldDef = Union[
     # Key, Value spec/annotation.

pyglove/core/typing/json_schema.py

@@ -15,7 +15,7 @@
 
 import dataclasses
 import inspect
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union
 
 from pyglove.core import utils
 from pyglove.core.typing import callable_signature
@@ -53,8 +53,8 @@ def _json_schema_from_schema(
     required = []
 
     if type_name and include_type_name:
-      properties['_type'] = {'const': type_name}
-      required.append('_type')
+      properties[utils.JSONConvertible.TYPE_NAME_KEY] = {'const': type_name}
+      required.append(utils.JSONConvertible.TYPE_NAME_KEY)
 
     for key, field in schema.items():
       if isinstance(key, ks.ConstStrKey):
@@ -145,7 +145,6 @@ def _json_schema_from_value_spec(
     }
     if not isinstance(value_spec.element.value, vs.Any):
       definition['items'] = _child_json_schema(value_spec.element.value)
-    return definition
   elif isinstance(value_spec, vs.Dict):
     if value_spec.schema is None:
       definition = {
@@ -168,9 +167,7 @@ def _json_schema_from_value_spec(
           include_type_name=include_type_name,
           include_subclasses=include_subclasses,
       )
-    definitions = [
-        _json_schema_from_cls(value_spec.cls)
-    ]
+    definitions = [_json_schema_from_cls(value_spec.cls)]
 
     if include_subclasses:
       for subclass in value_spec.cls.__subclasses__():
@@ -206,12 +203,26 @@ def _json_schema_from_value_spec(
         f'Value spec {value_spec!r} cannot be converted to JSON schema.'
     )
 
+  if (value_spec.has_default
+      and value_spec.default is not None
+      and not isinstance(value_spec, vs.Dict)):
+    default = utils.to_json(value_spec.default)
+    if not include_type_name:
+      def _remove_type_name(_, v: Dict[str, Any]):
+        if isinstance(v, dict):
+          v.pop(utils.JSONConvertible.TYPE_NAME_KEY, None)
+        return v
+      default = utils.transform(default, _remove_type_name)
+    definition['default'] = default
+
   if not ignore_nonable and value_spec.is_noneable:
     nullable = {'type': 'null'}
     if 'anyOf' in definition:
       definition['anyOf'].append(nullable)
     else:
       definition = {'anyOf': [definition, nullable]}
+    if value_spec.default is None:
+      definition['default'] = None
   return definition
 
 
@@ -294,6 +305,8 @@ def _canonicalize_schema(
   }
   canonical_form = {'$defs': referenced_defs} if referenced_defs else {}
   canonical_form.update(new_root)
+  if 'default' in root:
+    canonical_form['default'] = root['default']
   return canonical_form
 
 #
@@ -336,5 +349,205 @@ def _schema_to_json_schema(
       **kwargs
   )
 
-vs.ValueSpec.to_json_schema = _value_spec_to_json_schema
+class_schema.ValueSpec.to_json_schema = _value_spec_to_json_schema
 class_schema.Schema.to_json_schema = _schema_to_json_schema
+
+#
+# from JSON schema to PyGlove value spec.
+#
+
+
+def _json_schema_to_value_spec(
+    json_schema: Dict[str, Any],
+    defs: Dict[str, Type[Any]],
+    class_fn: Optional[Callable[[str, class_schema.Schema], Type[Any]]],
+    add_json_schema_as_metadata: bool,
+) -> class_schema.ValueSpec:
+  """Converts a JSON schema to a value spec."""
+  # Generate code to convert JSON schema to value spec.
+  def _value_spec(value_schema: Dict[str, Any]):
+    return _json_schema_to_value_spec(
+        value_schema, defs, class_fn, add_json_schema_as_metadata
+    )
+
+  if '$ref' in json_schema:
+    # TODO(daiyip): Support circular references.
+    ref_key = json_schema['$ref'].split('/')[-1]
+    type_ref = defs.get(ref_key)
+    if type_ref is None:
+      raise ValueError(
+          f'Reference {ref_key!r} not defined in defs. '
+          'Please make sure classes being referenced are defined '
+          'before the referencing classes. '
+      )
+    return type_ref
+  type_str = json_schema.get('type')
+  default = json_schema.get('default', utils.MISSING_VALUE)
+  if type_str is None:
+    if 'enum' in json_schema:
+      for v in json_schema['enum']:
+        if not isinstance(v, (str, int, float, bool)):
+          raise ValueError(
+              f'Enum candidate {v!r} is not supported for JSON schema '
+              'conversion.'
+          )
+      return vs.Enum(
+          default,
+          [v for v in json_schema['enum'] if v is not None]
+      )
+    elif 'anyOf' in json_schema:
+      candidates = []
+      accepts_none = False
+      for v in json_schema['anyOf']:
+        candidate = _value_spec(v)
+        if candidate.frozen and candidate.default is None:
+          accepts_none = True
+          continue
+        candidates.append(candidate)
+
+      if len(candidates) == 1:
+        spec = candidates[0]
+      else:
+        spec = vs.Union(candidates)
+      if accepts_none:
+        spec = spec.noneable()
+      return spec
+  elif type_str == 'null':
+    return vs.Any().freeze(None)
+  elif type_str == 'boolean':
+    return vs.Bool(default=default)
+  elif type_str == 'integer':
+    minimum = json_schema.get('minimum')
+    maximum = json_schema.get('maximum')
+    return vs.Int(min_value=minimum, max_value=maximum, default=default)
+  elif type_str == 'number':
+    minimum = json_schema.get('minimum')
+    maximum = json_schema.get('maximum')
+    return vs.Float(min_value=minimum, max_value=maximum, default=default)
+  elif type_str == 'string':
+    pattern = json_schema.get('pattern')
+    return vs.Str(regex=pattern, default=default)
+  elif type_str == 'array':
+    items = json_schema.get('items')
+    return vs.List(_value_spec(items) if items else vs.Any(), default=default)
+  elif type_str == 'object':
+    schema = _json_schema_to_schema(
+        json_schema, defs, class_fn, add_json_schema_as_metadata
+    )
+    if class_fn is not None and 'title' in json_schema:
+      return vs.Object(class_fn(json_schema['title'], schema))
+    return vs.Dict(schema=schema if schema.fields else None)
+  raise ValueError(f'Unsupported type {type_str!r} in JSON schema.')
+
+
+def _json_schema_to_schema(
+    json_schema: Dict[str, Any],
+    defs: Dict[str, Type[Any]],
+    class_fn: Optional[Callable[[str, class_schema.Schema], Type[Any]]],
+    add_json_schema_as_metadata: bool,
+) -> class_schema.Schema:
+  """Converts a JSON schema to a schema."""
+  title = json_schema.get('title')
+  properties = json_schema.get('properties', {})
+  fields = []
+  required = set(json_schema.get('required', []))
+  for name, property_schema in properties.items():
+    value_spec = _json_schema_to_value_spec(
+        property_schema, defs, class_fn, add_json_schema_as_metadata
+    )
+    if name not in required and not value_spec.has_default:
+      value_spec = value_spec.noneable()
+    fields.append(
+        class_schema.Field(
+            name,
+            value_spec,
+            description=property_schema.get('description'),
+            metadata=(
+                dict(json_schema=property_schema)
+                if add_json_schema_as_metadata else None
+            )
+        )
+    )
+  additional_properties = json_schema.get('additionalProperties')
+  if additional_properties:
+    if isinstance(additional_properties, dict):
+      value_spec = _json_schema_to_value_spec(
+          additional_properties, defs, class_fn, add_json_schema_as_metadata
+      )
+    else:
+      value_spec = vs.Any()
+    fields.append(class_schema.Field(ks.StrKey(), value_spec))
+  return class_schema.Schema(
+      name=title,
+      description=json_schema.get('description'),
+      fields=fields,
+      allow_nonconst_keys=True,
+  )
+
+
+@classmethod
+def _value_spec_from_json_schema(
+    cls,
+    json_schema: Dict[str, Any],
+    class_fn: Optional[Callable[[str, class_schema.Schema], Type[Any]]] = None,
+    add_json_schema_as_metadata: bool = False,
+) -> class_schema.ValueSpec:
+  """Creates a PyGlove value spec from a JSON schema.
+
+  Args:
+    json_schema: The JSON schema for a value spec.
+    class_fn: A function that creates a PyGlove class from a class name and a
+      schema. If None, all "object" type properties will be converted to
+      `pg.typing.Dict`. Otherwise, "object" type properties will be converted to
+      a class.
+    add_json_schema_as_metadata: Whether to add the JSON schema as field
+      metadata.
+
+  Returns:
+    A PyGlove value spec.
+  """
+  del cls
+  defs = {}
+  if '$defs' in json_schema:
+    for key, def_entry in json_schema['$defs'].items():
+      defs[key] = _json_schema_to_value_spec(
+          def_entry, defs, class_fn, add_json_schema_as_metadata
+      )
+  return _json_schema_to_value_spec(
+      json_schema, defs, class_fn, add_json_schema_as_metadata
+  )
+
+
+@classmethod
+def _schema_from_json_schema(
+    cls,
+    json_schema: Dict[str, Any],
+    class_fn: Optional[Callable[[str, class_schema.Schema], Type[Any]]] = None,
+    add_json_schema_as_metadata: bool = False,
+) -> class_schema.Schema:
+  """Creates a PyGlove schema from a JSON schema.
+
+  Args:
+    json_schema: The JSON schema to convert.
+    class_fn: A function that creates a PyGlove class from a class name and a
+      schema. If None, all "object" type properties will be converted to
+      `pg.typing.Dict`. Otherwise, "object" type properties will be converted to
+      a class.
+    add_json_schema_as_metadata: Whether to add the JSON schema as field
+      metadata.
+
+  Returns:
+    A PyGlove schema.
+  """
+  del cls
+  if json_schema.get('type') != 'object':
+    raise ValueError(
+        f'JSON schema is not an object type: {json_schema!r}'
+    )
+  return _json_schema_to_schema(
+      json_schema, {}, class_fn, add_json_schema_as_metadata
+  )
+
+
+class_schema.ValueSpec.from_json_schema = _value_spec_from_json_schema
+class_schema.Schema.from_json_schema = _schema_from_json_schema

pyglove/core/typing/json_schema_test.py

@@ -31,7 +31,7 @@ class Bar(pg_object.Object):
   z: Optional[Foo]
 
 
-class JsonSchemaTest(unittest.TestCase):
+class ToJsonSchemaTest(unittest.TestCase):
 
   maxDiff = None
 
@@ -53,6 +53,10 @@ class JsonSchemaTest(unittest.TestCase):
     self.assert_json_schema(bool, {
         'type': 'boolean',
     })
+    self.assert_json_schema(vs.Bool(default=True), {
+        'type': 'boolean',
+        'default': True,
+    })
 
   def test_int(self):
     self.assert_json_schema(int, {
@@ -62,9 +66,10 @@ class JsonSchemaTest(unittest.TestCase):
         'type': 'integer',
         'minimum': 0,
     })
-    self.assert_json_schema(vs.Int(max_value=1), {
+    self.assert_json_schema(vs.Int(max_value=1, default=0), {
         'type': 'integer',
         'maximum': 1,
+        'default': 0,
     })
 
   def test_float(self):
@@ -75,30 +80,43 @@ class JsonSchemaTest(unittest.TestCase):
         'type': 'number',
         'minimum': 0.0,
     })
-    self.assert_json_schema(vs.Float(max_value=1.0), {
+    self.assert_json_schema(vs.Float(max_value=1.0, default=0.0), {
         'type': 'number',
         'maximum': 1.0,
+        'default': 0.0,
     })
 
   def test_str(self):
     self.assert_json_schema(str, {
         'type': 'string',
     })
-    self.assert_json_schema(vs.Str(regex='a.*'), {
+    self.assert_json_schema(vs.Str(regex='a.*', default='a1'), {
         'type': 'string',
         'pattern': 'a.*',
+        'default': 'a1',
     })
 
   def test_enum(self):
     self.assert_json_schema(Literal['a', 1], {
         'enum': ['a', 1]
     })
+    self.assert_json_schema(vs.Enum(1, ['a', 1]), {
+        'enum': ['a', 1],
+        'default': 1,
+    })
     self.assert_json_schema(Literal['a', 1, None], {
         'anyOf': [
             {'enum': ['a', 1]},
             {'type': 'null'}
         ]
     })
+    self.assert_json_schema(vs.Enum(None, ['a', 1, None]), {
+        'anyOf': [
+            {'enum': ['a', 1]},
+            {'type': 'null'}
+        ],
+        'default': None
+    })
     with self.assertRaisesRegex(
         ValueError, 'Enum candidate .* is not supported'
     ):
@@ -111,6 +129,13 @@ class JsonSchemaTest(unittest.TestCase):
             'type': 'integer',
         }
     })
+    self.assert_json_schema(vs.List(int, default=[1, 2]), {
+        'type': 'array',
+        'items': {
+            'type': 'integer',
+        },
+        'default': [1, 2],
+    })
 
   def test_dict(self):
     self.assert_json_schema(vs.Dict(), {
@@ -143,6 +168,15 @@ class JsonSchemaTest(unittest.TestCase):
             {'type': 'null'},
         ]
     })
+    self.assert_json_schema(vs.Union([int, vs.Union([str, int]).noneable()]), {
+        'anyOf': [
+            {'type': 'integer'},
+            {'type': 'string'},
+            # TODO(daiyip): Remove duplicates for nested Union in future.
+            {'type': 'integer'},
+            {'type': 'null'},
+        ]
+    })
 
   def test_any(self):
     self.assert_json_schema(vs.Any(), {
@@ -155,6 +189,17 @@ class JsonSchemaTest(unittest.TestCase):
             {'type': 'null'},
         ]
     })
+    self.assert_json_schema(vs.Any(default=1), {
+        'anyOf': [
+            {'type': 'boolean'},
+            {'type': 'number'},
+            {'type': 'string'},
+            {'type': 'array'},
+            {'type': 'object', 'additionalProperties': True},
+            {'type': 'null'},
+        ],
+        'default': 1,
+    })
 
   def test_object(self):
     class A:
@@ -176,9 +221,16 @@ class JsonSchemaTest(unittest.TestCase):
         'additionalProperties': False,
     }, include_type_name=False)
 
-    self.assert_json_schema(vs.Object(A), {
+    class B(pg_object.Object):
+      x: int
+      y: str
+
+    self.assert_json_schema(vs.Object(B), {
         'type': 'object',
         'properties': {
+            '_type': {
+                'const': B.__type_name__,
+            },
             'x': {
                 'type': 'integer',
             },
@@ -186,11 +238,30 @@ class JsonSchemaTest(unittest.TestCase):
                 'type': 'string',
             },
         },
-        'required': ['x', 'y'],
-        'title': 'A',
+        'required': ['_type', 'x', 'y'],
+        'title': 'B',
         'additionalProperties': False,
     }, include_type_name=True)
 
+    self.assert_json_schema(vs.Object(B, default=B(x=1, y='a')), {
+        'type': 'object',
+        'properties': {
+            'x': {
+                'type': 'integer',
+            },
+            'y': {
+                'type': 'string',
+            },
+        },
+        'required': ['x', 'y'],
+        'title': 'B',
+        'additionalProperties': False,
+        'default': {
+            'x': 1,
+            'y': 'a',
+        },
+    }, include_type_name=False)
+
   def test_pg_object(self):
 
     class A(pg_object.Object):
@@ -230,14 +301,14 @@ class JsonSchemaTest(unittest.TestCase):
         'additionalProperties': False,
     }, include_type_name=True)
 
-  def test_pg_object_nessted(self):
+  def test_pg_object_nested(self):
 
     class A(pg_object.Object):
       x: Annotated[int, 'field x']
       y: str
 
     class B(pg_object.Object):
-      z: A
+      z: A = A(x=1, y='a')
 
     self.assert_json_schema(vs.Object(B), {
         '$defs': {
@@ -266,10 +337,15 @@ class JsonSchemaTest(unittest.TestCase):
                 'const': B.__type_name__,
             },
             'z': {
-                '$ref': '#/$defs/A'
+                '$ref': '#/$defs/A',
+                'default': {
+                    '_type': A.__type_name__,
+                    'x': 1,
+                    'y': 'a',
+                },
             },
         },
-        'required': ['_type', 'z'],
+        'required': ['_type'],
         'title': 'B',
         'additionalProperties': False,
     }, include_type_name=True)
@@ -299,7 +375,7 @@ class JsonSchemaTest(unittest.TestCase):
                 'additionalProperties': False,
             },
         },
-        'required': ['_type', 'z'],
+        'required': ['_type'],
         'title': 'B',
         'additionalProperties': False,
     }, include_type_name=True, inline_nested_refs=True)
@@ -357,6 +433,7 @@ class JsonSchemaTest(unittest.TestCase):
                     'type': 'null',
                 }
             ],
+            'default': None,
         },
         include_type_name=True,
         include_subclasses=True,
@@ -473,5 +550,424 @@ class JsonSchemaTest(unittest.TestCase):
         }
     )
 
+
+class FromJsonSchemaTest(unittest.TestCase):
+
+  def assert_value_spec(self, input_json_schema, expected_value_spec):
+    value_spec = vs.ValueSpec.from_json_schema(input_json_schema)
+    self.assertEqual(value_spec, expected_value_spec)
+
+  def test_bool(self):
+    self.assert_value_spec(
+        {
+            'type': 'boolean',
+        },
+        vs.Bool(),
+    )
+    self.assert_value_spec(
+        {
+            'type': 'boolean',
+            'default': True
+        },
+        vs.Bool(default=True),
+    )
+
+  def test_int(self):
+    self.assert_value_spec(
+        {
+            'type': 'integer',
+        },
+        vs.Int(),
+    )
+    self.assert_value_spec(
+        {
+            'type': 'integer',
+            'minimum': 0,
+        },
+        vs.Int(min_value=0),
+    )
+    self.assert_value_spec(
+        {
+            'type': 'integer',
+            'maximum': 1,
+            'default': 0,
+        },
+        vs.Int(max_value=1, default=0),
+    )
+
+  def test_number(self):
+    self.assert_value_spec(
+        {
+            'type': 'number',
+        },
+        vs.Float(),
+    )
+    self.assert_value_spec(
+        {
+            'type': 'number',
+            'minimum': 0.0,
+        },
+        vs.Float(min_value=0.0),
+    )
+    self.assert_value_spec(
+        {
+            'type': 'number',
+            'maximum': 1.0,
+            'default': 0.0,
+        },
+        vs.Float(max_value=1.0, default=0.0),
+    )
+
+  def test_str(self):
+    self.assert_value_spec(
+        {
+            'type': 'string',
+        },
+        vs.Str(),
+    )
+    self.assert_value_spec(
+        {
+            'type': 'string',
+            'pattern': 'a.*',
+            'default': 'a1',
+        },
+        vs.Str(regex='a.*', default='a1'),
+    )
+
+  def test_enum(self):
+    self.assert_value_spec(
+        {
+            'enum': ['a', 'b', 'c'],
+            'default': 'b',
+        },
+        vs.Enum('b', ['a', 'b', 'c']),
+    )
+    with self.assertRaisesRegex(
+        ValueError, 'Enum candidate .* is not supported'
+    ):
+      vs.ValueSpec.from_json_schema({'enum': [{'x': 1}, {'y': 'abc'}]})
+
+  def test_null(self):
+    self.assert_value_spec(
+        {
+            'type': 'null',
+        },
+        vs.Any().freeze(None),
+    )
+
+  def test_any_of(self):
+    self.assert_value_spec(
+        {
+            'anyOf': [
+                {'type': 'integer'},
+            ],
+        },
+        vs.Int(),
+    )
+    self.assert_value_spec(
+        {
+            'anyOf': [
+                {'type': 'integer'},
+                {'type': 'string'},
+            ],
+        },
+        vs.Union([vs.Int(), vs.Str()]),
+    )
+    self.assert_value_spec(
+        {
+            'anyOf': [
+                {'type': 'integer'},
+                {'type': 'string'},
+                {'type': 'null'},
+            ],
+        },
+        vs.Union([vs.Int(), vs.Str()]).noneable(),
+    )
+
+  def test_list(self):
+    self.assert_value_spec(
+        {
+            'type': 'array',
+        },
+        vs.List(vs.Any()),
+    )
+    self.assert_value_spec(
+        {
+            'type': 'array',
+            'items': {
+                'type': 'integer',
+            },
+        },
+        vs.List(vs.Int()),
+    )
+    self.assert_value_spec(
+        {
+            'type': 'array',
+            'items': {
+                'type': 'integer',
+            },
+            'default': [1, 2],
+        },
+        vs.List(vs.Int(), default=[1, 2]),
+    )
+
+  def test_dict(self):
+    self.assert_value_spec(
+        {
+            'type': 'object',
+        },
+        vs.Dict(),
+    )
+    self.assert_value_spec(
+        {
+            'type': 'object',
+            'properties': {
+                'a': {
+                    'type': 'integer',
+                },
+            },
+            'required': ['a'],
+            'additionalProperties': False,
+        },
+        vs.Dict({'a': vs.Int()}),
+    )
+    self.assert_value_spec(
+        {
+            'type': 'object',
+            'additionalProperties': {'type': 'integer'},
+        },
+        vs.Dict([(ks.StrKey(), vs.Int())]),
+    )
+    self.assert_value_spec(
+        {
+            'type': 'object',
+            'additionalProperties': True,
+        },
+        vs.Dict([(ks.StrKey(), vs.Any())]),
+    )
+
+  def _cls_value_spec(self, input_json_schema):
+    def schema_to_class(name, schema):
+      class _Class(pg_object.Object):
+        pass
+      cls = _Class
+      cls.__name__ = name
+      cls.__doc__ = schema.description
+      cls.apply_schema(schema)
+      return cls
+    return vs.ValueSpec.from_json_schema(
+        input_json_schema, class_fn=schema_to_class
+    )
+
+  def test_simple_object(self):
+    cls_spec = self._cls_value_spec(
+        {
+            'type': 'object',
+            'title': 'A',
+            'description': 'Class A',
+            'properties': {
+                'x': {
+                    'type': 'integer',
+                    'description': 'field x',
+                },
+                'y': {
+                    'type': 'string',
+                },
+            },
+            'required': ['x'],
+            'additionalProperties': False,
+        },
+    )
+    self.assertIsNone(cls_spec.cls(x=1).y)
+    self.assertEqual(cls_spec.cls.__name__, 'A')
+    self.assertEqual(cls_spec.cls.__doc__, 'Class A')
+    self.assertEqual(
+        cls_spec.cls.__schema__['x'], vs.Field('x', vs.Int(), 'field x')
+    )
+    self.assertEqual(
+        cls_spec.cls.__schema__['y'], vs.Field('y', vs.Str().noneable())
+    )
+
+  def test_nested_object(self):
+    cls_spec = self._cls_value_spec(
+        {
+            'type': 'object',
+            'title': 'A',
+            'description': 'Class A',
+            'properties': {
+                'x': {
+                    'type': 'integer',
+                    'description': 'field x',
+                },
+                'y': {
+                    'type': 'object',
+                    'title': 'B',
+                    'description': 'Class B',
+                    'properties': {
+                        'z': {
+                            'type': 'string',
+                        },
+                    },
+                    'required': ['z'],
+                    'additionalProperties': False,
+                },
+            },
+            'required': ['x'],
+            'additionalProperties': False,
+        },
+    )
+    self.assertIsNone(cls_spec.cls(x=1).y)
+    self.assertEqual(cls_spec.cls.__name__, 'A')
+    self.assertEqual(cls_spec.cls.__doc__, 'Class A')
+    self.assertEqual(
+        cls_spec.cls.__schema__['x'], vs.Field('x', vs.Int(), 'field x')
+    )
+    b_cls = cls_spec.cls.__schema__['y'].value.cls
+    self.assertEqual(b_cls.__schema__['z'], vs.Field('z', vs.Str()))
+
+  def test_simple_object_with_def(self):
+    cls_spec = self._cls_value_spec(
+        {
+            '$defs': {
+                'A': {
+                    'type': 'object',
+                    'title': 'A',
+                    'description': 'Class A',
+                    'properties': {
+                        'x': {
+                            'type': 'integer',
+                            'description': 'field x',
+                            'default': 1,
+                        },
+                        'y': {
+                            'type': 'string',
+                        },
+                    },
+                    'required': ['x'],
+                    'additionalProperties': False,
+                }
+            },
+            '$ref': '#/$defs/A',
+        }
+    )
+    self.assertEqual(cls_spec.cls(y='a').x, 1)
+    self.assertEqual(cls_spec.cls.__name__, 'A')
+    self.assertEqual(cls_spec.cls.__doc__, 'Class A')
+
+  def test_complex_object_with_def(self):
+    cls_spec = self._cls_value_spec(
+        {
+            '$defs': {
+                'B': {
+                    'type': 'object',
+                    'title': 'B',
+                    'description': 'Class B',
+                    'properties': {
+                        'z': {
+                            'type': 'string',
+                        },
+                    },
+                    'required': ['z'],
+                    'additionalProperties': False,
+                },
+                'A': {
+                    'type': 'object',
+                    'title': 'A',
+                    'description': 'Class A',
+                    'properties': {
+                        'x': {
+                            'type': 'integer',
+                            'description': 'field x',
+                            'default': 1,
+                        },
+                        'y': {
+                            '$ref': '#/$defs/B',
+                        }
+                    },
+                    'required': ['x'],
+                    'additionalProperties': False,
+                },
+            },
+            '$ref': '#/$defs/A',
+        }
+    )
+    self.assertIsNone(cls_spec.cls(x=1).y)
+    self.assertEqual(cls_spec.cls.__name__, 'A')
+    self.assertEqual(cls_spec.cls.__doc__, 'Class A')
+    self.assertEqual(
+        cls_spec.cls.__schema__['x'],
+        vs.Field('x', vs.Int(default=1), 'field x')
+    )
+    b_cls = cls_spec.cls.__schema__['y'].value.cls
+    self.assertEqual(b_cls.__name__, 'B')
+    self.assertEqual(b_cls.__doc__, 'Class B')
+    self.assertEqual(b_cls.__schema__['z'], vs.Field('z', vs.Str()))
+
+    with self.assertRaisesRegex(
+        ValueError, 'Reference .* not defined'
+    ):
+      self._cls_value_spec(
+          {
+              '$defs': {
+                  'A': {
+                      'type': 'object',
+                      'title': 'A',
+                      'description': 'Class A',
+                      'properties': {
+                          'x': {
+                              '$ref': '#/$defs/B',
+                          }
+                      },
+                      'required': ['x'],
+                      'additionalProperties': False,
+                  },
+                  # B should go before A.
+                  'B': {
+                      'type': 'object',
+                      'title': 'B',
+                      'description': 'Class B',
+                      'properties': {
+                          'z': {
+                              'type': 'string',
+                          },
+                      },
+                      'required': ['z'],
+                      'additionalProperties': False,
+                  },
+              },
+              '$ref': '#/$defs/A',
+          }
+      )
+
+  def test_unsupported_json_schema(self):
+    with self.assertRaisesRegex(
+        ValueError, 'Unsupported type .* in JSON schema'
+    ):
+      vs.ValueSpec.from_json_schema({'type': 'oneOf'})
+
+  def test_schema_from_json_schema(self):
+    schema = vs.Schema.from_json_schema(
+        {
+            'type': 'object',
+            'title': 'A',
+            'description': 'Class A',
+            'properties': {
+                'x': {
+                    'type': 'integer',
+                },
+            },
+            'required': ['x'],
+            'additionalProperties': False,
+        },
+    )
+    self.assertEqual(schema.description, 'Class A')
+    self.assertEqual(list(schema.fields.keys()), ['x'])
+    self.assertEqual(schema.fields['x'].value, vs.Int())
+
+    with self.assertRaisesRegex(
+        ValueError, 'JSON schema is not an object type'
+    ):
+      vs.Schema.from_json_schema({'type': 'integer'})
+
 if __name__ == '__main__':
   unittest.main()

pyglove/core/utils/contextual.py

@@ -15,6 +15,7 @@
 
 import contextlib
 import dataclasses
+import inspect
 import threading
 from typing import Any, Callable, ContextManager, Iterator, Optional
 
@@ -89,10 +90,14 @@ def with_contextual_override(func: Callable[..., Any]) -> Callable[..., Any]:
   with contextual_override() as current_context:
     pass
 
-  def _func(*args, **kwargs) -> Any:
-    with contextual_override(**current_context):
-      return func(*args, **kwargs)
-
+  if inspect.iscoroutinefunction(func):
+    async def _func(*args, **kwargs) -> Any:
+      with contextual_override(**current_context):
+        return await func(*args, **kwargs)
+  else:
+    def _func(*args, **kwargs) -> Any:
+      with contextual_override(**current_context):
+        return func(*args, **kwargs)
   return _func
 
 

pyglove/core/utils/contextual_test.py

@@ -11,6 +11,7 @@
 # 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.
+import asyncio
 import concurrent.futures
 import unittest
 from pyglove.core.utils import contextual
@@ -83,6 +84,15 @@ class ContextualTest(unittest.TestCase):
           [3]
       )
 
+  def test_with_contextual_override_async_func(self):
+    async def func(i):
+      del i
+      return contextual.contextual_value('x')
+
+    with contextual.contextual_override(x=3):
+      self.assertEqual(
+          asyncio.run(contextual.with_contextual_override(func)(0)), 3
+      )
 
 if __name__ == '__main__':
   unittest.main()

{pyglove-0.5.0.dev202510150810.dist-info → pyglove-0.5.0.dev202510170226.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pyglove
-Version: 0.5.0.dev202510150810
+Version: 0.5.0.dev202510170226
 Summary: PyGlove: A library for manipulating Python objects.
 Home-page: https://github.com/google/pyglove
 Author: PyGlove Authors

{pyglove-0.5.0.dev202510150810.dist-info → pyglove-0.5.0.dev202510170226.dist-info}/RECORD

@@ -120,13 +120,13 @@ pyglove/core/typing/callable_ext.py,sha256=PiBQWPeUAH7Lgmf2xKCZqgK7N0OSrTdbnEkV8
 pyglove/core/typing/callable_ext_test.py,sha256=TnWKU4_ZjvpbHZFtFHgFvCMDiCos8VmLlODcM_7Xg8M,10156
 pyglove/core/typing/callable_signature.py,sha256=DRpt7aShfkn8pb3SCiZzS_27eHbkQ_d2UB8BUhJjs0Q,27176
 pyglove/core/typing/callable_signature_test.py,sha256=iQmHsKPhJPQlMikDhEyxKyq7yWyXI9juKCLYgKhrH3U,25145
-pyglove/core/typing/class_schema.py,sha256=3U92-Suwdxlc8mXjK6Tqa9iiqogtRuv1YZwVYpmkMTc,55288
+pyglove/core/typing/class_schema.py,sha256=jICxKo6Ubu35FJO1ei6iOqHB9I204Wtea5UGng2_hTE,55897
 pyglove/core/typing/class_schema_test.py,sha256=sJkE7ndDSIKb0EUcjZiVFOeJYDI7Hdu2GdPJCMgZxrI,29556
 pyglove/core/typing/custom_typing.py,sha256=qdnIKHWNt5kZAAFdpQXra8bBu6RljMbbJ_YDG2mhAUA,2205
 pyglove/core/typing/inspect.py,sha256=VLSz1KAunNm2hx0eEMjiwxKLl9FHlKr9nHelLT25iEA,7726
 pyglove/core/typing/inspect_test.py,sha256=xclevobF0X8c_B5b1q1dkBJZN1TsVA1RUhk5l25DUCM,10248
-pyglove/core/typing/json_schema.py,sha256=shHElSIpw6AZ5L1iKrKzX3yZXgGKYFJOvMC-syEf-i8,10453
-pyglove/core/typing/json_schema_test.py,sha256=DhtyoWLoYzCS-DU67QuwaZKZ-9WgK_ZFVsiFp94Ltx8,13011
+pyglove/core/typing/json_schema.py,sha256=F7uWXN8Tk4cAydh8FE9jncaP_n7KEvgprvNHeagTvqE,17871
+pyglove/core/typing/json_schema_test.py,sha256=ZxMO2xgKiELNDzoQ84cmXsyCtFA0Ltn1I4gZjF-3efY,26482
 pyglove/core/typing/key_specs.py,sha256=-7xjCuUGoQgD0sMafsRFNlw3S4f1r-7t5OO4ev5bbeI,9225
 pyglove/core/typing/key_specs_test.py,sha256=5zornpyHMAYoRaG8KDXHiQ3obu9UfRp3399lBeUNTPk,6499
 pyglove/core/typing/pytype_support.py,sha256=lyX11WVbCwoOi5tTQ90pEOS-yvo_6iEi7Lxbp-nXu2A,2069
@@ -139,8 +139,8 @@ pyglove/core/typing/value_specs_test.py,sha256=eGXVxdduIM-oEaapJS9Kh7WSQHRUFegLI
 pyglove/core/utils/__init__.py,sha256=2aw4n1kYG9xlX2tWI-H5i25cBuK1ME9Lmf-F31VlKEk,8657
 pyglove/core/utils/common_traits.py,sha256=PWxOgPhG5H60ZwfO8xNAEGRjFUqqDZQBWQYomOfvdy8,3640
 pyglove/core/utils/common_traits_test.py,sha256=DIuZB_1xfmeTVfWnGOguDQcDAM_iGgBOe8C-5CsIqBc,1122
-pyglove/core/utils/contextual.py,sha256=RxBQkDM2gB6QwZj_2oMels6oh-zQPGJJlinZbbqHuYQ,5148
-pyglove/core/utils/contextual_test.py,sha256=OOcthquVyAekTqt1RyYcEMHaockMIokpbv4pSf13Nzs,3252
+pyglove/core/utils/contextual.py,sha256=_EO_ubCcmI81QyYyyucm3QcH1asQWSeXvT2Xa4KONs8,5355
+pyglove/core/utils/contextual_test.py,sha256=g5P6IBf6XgQTUW83QTDC1K5JQJ9bxP6ULiAj2ko4mMo,3552
 pyglove/core/utils/docstr_utils.py,sha256=5BY40kXozPKVGOB0eN8jy1P5_GHIzqFJ9FXAu_kzxaw,5119
 pyglove/core/utils/docstr_utils_test.py,sha256=TBHNwvhGyyoEs7dNOv5bW7h3B_y2smDyoAR9uPDiidI,4179
 pyglove/core/utils/error_utils.py,sha256=5if1LlouHMx9EkZt-rhlkDU53Qc-hmXgjhJHhhI125s,6392
@@ -218,8 +218,8 @@ pyglove/ext/scalars/randoms.py,sha256=LkMIIx7lOq_lvJvVS3BrgWGuWl7Pi91-lA-O8x_gZs
 pyglove/ext/scalars/randoms_test.py,sha256=nEhiqarg8l_5EOucp59CYrpO2uKxS1pe0hmBdZUzRNM,2000
 pyglove/ext/scalars/step_wise.py,sha256=IDw3tuTpv0KVh7AN44W43zqm1-E0HWPUlytWOQC9w3Y,3789
 pyglove/ext/scalars/step_wise_test.py,sha256=TL1vJ19xVx2t5HKuyIzGoogF7N3Rm8YhLE6JF7i0iy8,2540
-pyglove-0.5.0.dev202510150810.dist-info/licenses/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
-pyglove-0.5.0.dev202510150810.dist-info/METADATA,sha256=y7MjVTU0UfkHnZagBiYe7UeG3l6vS9CexBo3oGB4Bh0,7089
-pyglove-0.5.0.dev202510150810.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-pyglove-0.5.0.dev202510150810.dist-info/top_level.txt,sha256=wITzJSKcj8GZUkbq-MvUQnFadkiuAv_qv5qQMw0fIow,8
-pyglove-0.5.0.dev202510150810.dist-info/RECORD,,
+pyglove-0.5.0.dev202510170226.dist-info/licenses/LICENSE,sha256=WNHhf_5RCaeuKWyq_K39vmp9F28LxKsB4SpomwSZ2L0,11357
+pyglove-0.5.0.dev202510170226.dist-info/METADATA,sha256=Hk2SsgTsNgeAfgxQYyDRdCRj2I9g7Tq9sUdwdQ3GuMM,7089
+pyglove-0.5.0.dev202510170226.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+pyglove-0.5.0.dev202510170226.dist-info/top_level.txt,sha256=wITzJSKcj8GZUkbq-MvUQnFadkiuAv_qv5qQMw0fIow,8
+pyglove-0.5.0.dev202510170226.dist-info/RECORD,,