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

pyglove/core/symbolic/functor.py

@@ -22,15 +22,14 @@ import types
 import typing
 from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Type, Union
 
-from pyglove.core import object_utils
 from pyglove.core import typing as pg_typing
+from pyglove.core import utils
 from pyglove.core.symbolic import base
 from pyglove.core.symbolic import flags
 from pyglove.core.symbolic import object as pg_object
-from pyglove.core.symbolic import schema_utils
 
 
-class Functor(pg_object.Object, object_utils.Functor):
+class Functor(pg_object.Object, utils.Functor):
   """Symbolic functions (Functors).
 
   A symbolic function is a symbolic class with a ``__call__`` method, whose
@@ -99,7 +98,7 @@ class Functor(pg_object.Object, object_utils.Functor):
     # Update the return value of subclassed functors.
     if cls.is_subclassed_functor:  # pylint: disable=using-constant-test
       private_call_signature = pg_typing.Signature.from_callable(
-          cls._call, auto_typing=True
+          cls._call, auto_typing=True, auto_doc=True,
       )
       if (
           len(private_call_signature.args) > 1
@@ -125,7 +124,7 @@ class Functor(pg_object.Object, object_utils.Functor):
     if not hasattr(cls, '__orig_init__'):
       setattr(cls, '__orig_init__', cls.__init__)
 
-    @object_utils.explicit_method_override
+    @utils.explicit_method_override
     @functools.wraps(pseudo_init)
     def _init(self, *args, **kwargs):
       self.__class__.__orig_init__(self, *args, **kwargs)
@@ -149,14 +148,15 @@ class Functor(pg_object.Object, object_utils.Functor):
       return instance()
     return instance
 
-  @object_utils.explicit_method_override
+  @utils.explicit_method_override
   def __init__(
       self,
       *args,
-      root_path: Optional[object_utils.KeyPath] = None,
+      root_path: Optional[utils.KeyPath] = None,
       override_args: bool = False,
       ignore_extra_args: bool = False,
-      **kwargs):
+      **kwargs,
+  ):
     """Constructor.
 
     Args:
@@ -183,8 +183,8 @@ class Functor(pg_object.Object, object_utils.Functor):
         varargs = list(args[len(signature.args) :])
         args = args[: len(signature.args)]
       else:
-        arg_phrase = object_utils.auto_plural(len(signature.args), 'argument')
-        was_phrase = object_utils.auto_plural(len(args), 'was', 'were')
+        arg_phrase = utils.auto_plural(len(signature.args), 'argument')
+        was_phrase = utils.auto_plural(len(args), 'was', 'were')
         raise TypeError(
             f'{signature.id}() takes {len(signature.args)} '
             f'positional {arg_phrase} but {len(args)} {was_phrase} given.'
@@ -258,8 +258,7 @@ class Functor(pg_object.Object, object_utils.Functor):
     # pylint: enable=protected-access
     return typing.cast(Functor, other)
 
-  def _on_change(
-      self, field_updates: Dict[object_utils.KeyPath, base.FieldUpdate]):
+  def _on_change(self, field_updates: Dict[utils.KeyPath, base.FieldUpdate]):
     """Custom handling field change to update bound args."""
     for relative_path, update in field_updates.items():
       assert relative_path
@@ -376,7 +375,7 @@ class Functor(pg_object.Object, object_utils.Functor):
     if (
         signature.return_value
         and flags.is_type_check_enabled()
-        and return_value != pg_typing.MISSING_VALUE
+        and pg_typing.MISSING_VALUE != return_value
     ):
       return_value = signature.return_value.apply(
           return_value, root_path=self.sym_path + 'returns'
@@ -407,8 +406,8 @@ class Functor(pg_object.Object, object_utils.Functor):
       if ignore_extra_args:
         args = args[: len(signature.args)]
       else:
-        arg_phrase = object_utils.auto_plural(len(signature.args), 'argument')
-        was_phrase = object_utils.auto_plural(len(args), 'was', 'were')
+        arg_phrase = utils.auto_plural(len(signature.args), 'argument')
+        was_phrase = utils.auto_plural(len(args), 'was', 'were')
         raise TypeError(
             f'{signature.id}() takes {len(signature.args)} '
             f'positional {arg_phrase} but {len(args)} {was_phrase} given.'
@@ -426,7 +425,7 @@ class Functor(pg_object.Object, object_utils.Functor):
       varargs = list(args[len(signature.args) :])
       if flags.is_type_check_enabled():
         varargs = [
-            signature.varargs.value_spec.apply(
+            signature.varargs.value_spec.element.value.apply(
                 v, root_path=self.sym_path + signature.varargs.name
             )
             for v in varargs
@@ -484,9 +483,10 @@ class Functor(pg_object.Object, object_utils.Functor):
         missing_required_arg_names.append(arg.name)
 
     if missing_required_arg_names:
-      arg_phrase = object_utils.auto_plural(
-          len(missing_required_arg_names), 'argument')
-      args_str = object_utils.comma_delimited_str(missing_required_arg_names)
+      arg_phrase = utils.auto_plural(
+          len(missing_required_arg_names), 'argument'
+      )
+      args_str = utils.comma_delimited_str(missing_required_arg_names)
       raise TypeError(
           f'{signature.id}() missing {len(missing_required_arg_names)} '
           f'required positional {arg_phrase}: {args_str}.'
@@ -571,10 +571,11 @@ def functor(
 
 def functor_class(
     func: types.FunctionType,
-    args: Optional[List[Union[
-        Tuple[Tuple[str, pg_typing.KeySpec], pg_typing.ValueSpec, str],
-        Tuple[Tuple[str, pg_typing.KeySpec], pg_typing.ValueSpec, str, Any]]]
-    ] = None,   # pylint: disable=bad-continuation
+    args: Union[
+        List[Union[pg_typing.Field, pg_typing.FieldDef]],
+        Dict[pg_typing.FieldKeyDef, pg_typing.FieldValueDef],
+        None
+    ] = None,
     returns: Optional[pg_typing.ValueSpec] = None,
     base_class: Optional[Type[Functor]] = None,
     *,
@@ -684,9 +685,11 @@ def functor_class(
   cls.auto_register = True
 
   # Apply function schema.
-  schema = schema_utils.function_schema(
-      func, args, returns, auto_doc=auto_doc, auto_typing=auto_typing)
-  cls.apply_schema(schema)
+  cls.apply_schema(
+      pg_typing.schema(
+          func, args, returns, auto_doc=auto_doc, auto_typing=auto_typing
+      )
+  )
 
   # Register functor class for deserialization if needed.
   if add_to_registry:

pyglove/core/symbolic/functor_test.py

@@ -11,15 +11,13 @@
 # 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.
-"""Tests for pyglove.Functor."""
-
 import inspect
 import io
 import typing
 import unittest
 
-from pyglove.core import object_utils
 from pyglove.core import typing as pg_typing
+from pyglove.core import utils
 from pyglove.core.symbolic import flags
 from pyglove.core.symbolic.base import from_json_str as pg_from_json_str
 from pyglove.core.symbolic.dict import Dict
@@ -31,7 +29,7 @@ from pyglove.core.symbolic.object import members as pg_members
 from pyglove.core.symbolic.object import Object
 
 
-MISSING_VALUE = object_utils.MISSING_VALUE
+MISSING_VALUE = utils.MISSING_VALUE
 
 
 class FunctorTest(unittest.TestCase):
@@ -108,19 +106,41 @@ class FunctorTest(unittest.TestCase):
     self.assertEqual(
         f.__signature__.args,
         [
-            pg_typing.Argument('a', pg_typing.Any()),
-            pg_typing.Argument('b', pg_typing.Any()),
+            pg_typing.Argument(
+                'a',
+                pg_typing.Argument.Kind.POSITIONAL_OR_KEYWORD,
+                pg_typing.Any()
+            ),
+            pg_typing.Argument(
+                'b',
+                pg_typing.Argument.Kind.POSITIONAL_OR_KEYWORD,
+                pg_typing.Any()
+            ),
         ],
     )
     self.assertEqual(
-        f.__signature__.varargs, pg_typing.Argument('args', pg_typing.Any())
+        f.__signature__.varargs,
+        pg_typing.Argument(
+            'args',
+            pg_typing.Argument.Kind.VAR_POSITIONAL,
+            pg_typing.List(pg_typing.Any())
+        )
     )
     self.assertEqual(
-        f.__signature__.varkw, pg_typing.Argument('kwargs', pg_typing.Any())
+        f.__signature__.varkw,
+        pg_typing.Argument(
+            'kwargs',
+            pg_typing.Argument.Kind.VAR_KEYWORD,
+            pg_typing.Dict(pg_typing.Any())
+        )
     )
     self.assertEqual(
         f.__signature__.kwonlyargs,
-        [pg_typing.Argument('c', pg_typing.Any(default=0))],
+        [pg_typing.Argument(
+            'c',
+            pg_typing.Argument.Kind.KEYWORD_ONLY,
+            pg_typing.Any(default=0)
+        )],
     )
     self.assertIsNone(f.__signature__.return_value, None)
     self.assertTrue(f.__signature__.has_varargs)
@@ -148,8 +168,16 @@ class FunctorTest(unittest.TestCase):
     self.assertEqual(
         f.__signature__.args,
         [
-            pg_typing.Argument('a', pg_typing.Int(default=1)),
-            pg_typing.Argument('b', pg_typing.Int(default=2)),
+            pg_typing.Argument(
+                'a',
+                pg_typing.Argument.Kind.POSITIONAL_OR_KEYWORD,
+                pg_typing.Int(default=1)
+            ),
+            pg_typing.Argument(
+                'b',
+                pg_typing.Argument.Kind.POSITIONAL_OR_KEYWORD,
+                pg_typing.Int(default=2)
+            ),
         ],
     )
     self.assertEqual(
@@ -191,17 +219,36 @@ class FunctorTest(unittest.TestCase):
         ],
     )
     self.assertEqual(
-        f.__signature__.args, [pg_typing.Argument('a', pg_typing.Int())]
+        f.__signature__.args,
+        [pg_typing.Argument(
+            'a',
+            pg_typing.Argument.Kind.POSITIONAL_OR_KEYWORD,
+            pg_typing.Int()
+        )]
     )
     self.assertEqual(
-        f.__signature__.varargs, pg_typing.Argument('args', pg_typing.Any())
+        f.__signature__.varargs,
+        pg_typing.Argument(
+            'args',
+            pg_typing.Argument.Kind.VAR_POSITIONAL,
+            pg_typing.List(pg_typing.Any())
+        )
     )
     self.assertEqual(
         f.__signature__.kwonlyargs,
-        [pg_typing.Argument('b', pg_typing.Int(default=2))],
+        [pg_typing.Argument(
+            'b',
+            pg_typing.Argument.Kind.KEYWORD_ONLY,
+            pg_typing.Int(default=2)
+        )],
     )
     self.assertEqual(
-        f.__signature__.varkw, pg_typing.Argument('kwargs', pg_typing.Any())
+        f.__signature__.varkw,
+        pg_typing.Argument(
+            'kwargs',
+            pg_typing.Argument.Kind.VAR_KEYWORD,
+            pg_typing.Dict(pg_typing.Any())
+        )
     )
     self.assertEqual(f.__signature__.return_value, pg_typing.Int())
 
@@ -253,9 +300,21 @@ class FunctorTest(unittest.TestCase):
     self.assertEqual(
         f.__signature__.args,
         [
-            pg_typing.Argument('a', pg_typing.Int()),
-            pg_typing.Argument('b', pg_typing.Any(default=1)),
-            pg_typing.Argument('c', pg_typing.Int(default=1)),
+            pg_typing.Argument(
+                'a',
+                pg_typing.Argument.Kind.POSITIONAL_OR_KEYWORD,
+                pg_typing.Int()
+            ),
+            pg_typing.Argument(
+                'b',
+                pg_typing.Argument.Kind.POSITIONAL_OR_KEYWORD,
+                pg_typing.Any(default=1)
+            ),
+            pg_typing.Argument(
+                'c',
+                pg_typing.Argument.Kind.POSITIONAL_OR_KEYWORD,
+                pg_typing.Int(default=1)
+            ),
         ],
     )
     self.assertFalse(f.__signature__.has_varargs)
@@ -287,20 +346,31 @@ class FunctorTest(unittest.TestCase):
       def _call(self) -> int:
         return self.x + self.y
 
-    print(Foo.__signature__)
     self.assertEqual(Foo.__signature__.name, '__call__')
     self.assertEqual(Foo.__signature__.qualname, Foo.__qualname__)
     self.assertEqual(Foo.__signature__.module_name, Foo.__module__)
     self.assertEqual(
         Foo.__signature__.args,
         [
-            pg_typing.Argument('x', pg_typing.Int()),
-            pg_typing.Argument('y', pg_typing.Int()),
+            pg_typing.Argument(
+                'x',
+                pg_typing.Argument.Kind.POSITIONAL_OR_KEYWORD,
+                pg_typing.Int()
+            ),
+            pg_typing.Argument(
+                'y',
+                pg_typing.Argument.Kind.POSITIONAL_OR_KEYWORD,
+                pg_typing.Int()
+            ),
         ],
     )
     self.assertEqual(
         Foo.__signature__.varkw,
-        pg_typing.Argument('kwargs', pg_typing.Any(annotation=typing.Any)),
+        pg_typing.Argument(
+            'kwargs',
+            pg_typing.Argument.Kind.VAR_KEYWORD,
+            pg_typing.Dict(pg_typing.Any(annotation=typing.Any)),
+        ),
     )
     self.assertEqual(Foo.__signature__.return_value, pg_typing.Int())
 
@@ -422,7 +492,10 @@ class FunctorTest(unittest.TestCase):
     f = pg_as_functor(lambda x: x)
     self.assertIsInstance(f, Functor)
     self.assertEqual(
-        f.__signature__.args, [pg_typing.Argument('x', pg_typing.Any())]
+        f.__signature__.args,
+        [pg_typing.Argument(
+            'x', pg_typing.Argument.Kind.POSITIONAL_OR_KEYWORD, pg_typing.Any()
+        )]
     )
     self.assertIsNone(f.__signature__.return_value)
     self.assertEqual(f(1), 1)
@@ -453,41 +526,20 @@ class FunctorTest(unittest.TestCase):
     # `pg.functor` decorator found extra symbolic argument.
     with self.assertRaisesRegex(
         ValueError,
-        'the value spec for positional wildcard argument .*'
-        'must be a `pg.typing.List` instance'):
+        'Variable positional argument .* should have a value of .*List'
+    ):
       @pg_functor([('args', pg_typing.Int())])
       def f3(*args):  # pylint: disable=unused-variable
         del args
 
-    # `pg.functor` decorator has multiple StrKey.
-    with self.assertRaisesRegex(
-        KeyError,
-        '.* multiple StrKey found in symbolic arguments declaration.'):
-      @pg_functor([
-          ('a', pg_typing.Int()),
-          (pg_typing.StrKey(), pg_typing.Any()),
-          (pg_typing.StrKey(), pg_typing.Any()),
-      ])  # pylint: disable=unused-variable
-      def f4(a, **kwargs):
-        del a, kwargs
-
-    with self.assertRaisesRegex(
-        KeyError, '.* multiple symbolic fields found for argument.'):
-      @pg_functor([
-          ('a', pg_typing.Int()),
-          ('a', pg_typing.Str()),
-      ])  # pylint: disable=unused-variable
-      def f5(a):
-        del a
-
     with self.assertRaisesRegex(
         ValueError,
-        '.* the default value .* of symbolic argument .* does not equal '
-        'to the default value .* specified at function signature'):
+        'The annotated default value .* is not equal'
+    ):
       @pg_functor([
           ('a', pg_typing.Int(default=2)),
       ])  # pylint: disable=unused-variable
-      def f6(a=1):
+      def f4(a=1):
         del a
 
     with self.assertRaisesRegex(
@@ -495,7 +547,7 @@ class FunctorTest(unittest.TestCase):
       @pg_functor([
           ('a', pg_typing.Int()),
       ], returns=pg_typing.Any(default=None))
-      def f7(a=1):  # pylint: disable=unused-variable
+      def f5(a=1):  # pylint: disable=unused-variable
         del a
 
   def test_bad_call(self):

pyglove/core/symbolic/inferred.py

@@ -14,8 +14,8 @@
 """Common inferential values."""
 
 from typing import Any, Tuple
-from pyglove.core import object_utils
 from pyglove.core import typing as pg_typing
+from pyglove.core import utils
 from pyglove.core.symbolic import base
 from pyglove.core.symbolic.object import Object
 
@@ -65,7 +65,7 @@ class ValueFromParentChain(InferredValue):
       if v == pg_typing.MISSING_VALUE:
         if parent is None:
           raise AttributeError(
-              object_utils.message_on_path(
+              utils.message_on_path(
                   (
                       f'`{self.inference_key}` is not found under its context '
                       '(along its symbolic parent chain).'