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

pyglove/core/typing/callable_ext_test.py

@@ -18,6 +18,261 @@ from pyglove.core.typing import annotation_conversion   # pylint: disable=unused
 from pyglove.core.typing import callable_ext
 
 
+class PresetArgValueTest(unittest.TestCase):
+  """Tests for typing.PresetArgValue."""
+
+  def test_basics(self):
+    v = callable_ext.PresetArgValue()
+    self.assertFalse(v.has_default)
+    v = callable_ext.PresetArgValue(default=1)
+    self.assertTrue(v.has_default)
+    self.assertEqual(v.default, 1)
+    self.assertEqual(repr(v), 'PresetArgValue(default=1)')
+    self.assertEqual(str(v), 'PresetArgValue(\n  default=1\n)')
+    self.assertEqual(
+        callable_ext.PresetArgValue(), callable_ext.PresetArgValue()
+    )
+    self.assertEqual(
+        callable_ext.PresetArgValue(1), callable_ext.PresetArgValue(1)
+    )
+    self.assertNotEqual(
+        callable_ext.PresetArgValue(), callable_ext.PresetArgValue(default=1)
+    )
+
+  def test_inspect(self):
+
+    def foo(
+        x=callable_ext.PresetArgValue(),
+        y=1,
+        *,
+        z=callable_ext.PresetArgValue(default=2)
+    ):
+      return x + y + z
+
+    self.assertEqual(
+        callable_ext.PresetArgValue.inspect(foo),
+        dict(
+            x=callable_ext.PresetArgValue(),
+            z=callable_ext.PresetArgValue(default=2)
+        )
+    )
+
+    def bar(x, y=1):
+      return x + y
+
+    self.assertEqual(
+        callable_ext.PresetArgValue.inspect(bar),
+        {}
+    )
+
+  def test_resolve_args(self):
+    # Both positional and keyword arguments are from preset.
+    self.assertEqual(
+        callable_ext.PresetArgValue.resolve_args(
+            call_args=[],
+            call_kwargs=dict(),
+            positional_arg_names=['x', 'y'],
+            arg_defaults={
+                'x': callable_ext.PresetArgValue(),
+                'z': callable_ext.PresetArgValue(default=2),
+            },
+            preset_kwargs=dict(x=1, y=2, z=3, w=4)
+        ),
+        ([1], dict(z=3))
+    )
+    # Both positional and keyword arguments are absent and use the default.
+    self.assertEqual(
+        callable_ext.PresetArgValue.resolve_args(
+            call_args=[],
+            call_kwargs=dict(),
+            positional_arg_names=['x', 'y'],
+            arg_defaults={
+                'x': callable_ext.PresetArgValue(default=1),
+                'y': 0,
+                'z': callable_ext.PresetArgValue(default=2),
+            },
+            preset_kwargs=dict()
+        ),
+        ([1, 0], dict(z=2))
+    )
+    # Positional args from preset, keyword argument use preset default.
+    self.assertEqual(
+        callable_ext.PresetArgValue.resolve_args(
+            call_args=[],
+            call_kwargs=dict(),
+            positional_arg_names=['x', 'y'],
+            arg_defaults={
+                'x': callable_ext.PresetArgValue(),
+                'y': 0,
+                'z': callable_ext.PresetArgValue(default=2),
+            },
+            preset_kwargs=dict(x=1, y=2)
+        ),
+        ([1, 0], dict(z=2))
+    )
+    # Postional argument provided by user, which should take precedence over
+    # the preset value.
+    self.assertEqual(
+        callable_ext.PresetArgValue.resolve_args(
+            call_args=[2],
+            call_kwargs={},
+            positional_arg_names=['x', 'y'],
+            arg_defaults={
+                'x': callable_ext.PresetArgValue(),
+                'y': 0,
+                'z': callable_ext.PresetArgValue(default=2),
+            },
+            preset_kwargs=dict(x=1, y=2, z=3, w=4)
+        ),
+        ([2, 0], dict(z=3))
+    )
+    # Postional argument provided in keyword.
+    self.assertEqual(
+        callable_ext.PresetArgValue.resolve_args(
+            call_args=[],
+            call_kwargs=dict(x=2),
+            positional_arg_names=['x', 'y'],
+            arg_defaults={
+                'x': callable_ext.PresetArgValue(),
+                'y': 0,
+                'z': callable_ext.PresetArgValue(default=2),
+            },
+            preset_kwargs=dict(x=1, y=2, z=3, w=4)
+        ),
+        ([2, 0], dict(z=3))
+    )
+    # Postional argument provided in keyword, and there are more args
+    # (due to varargs)
+    self.assertEqual(
+        callable_ext.PresetArgValue.resolve_args(
+            call_args=[1, 2, 3],
+            call_kwargs=dict(x=2),
+            positional_arg_names=['x', 'y'],
+            arg_defaults={
+                'x': callable_ext.PresetArgValue(),
+                'y': 0,
+                'z': callable_ext.PresetArgValue(default=2),
+            },
+            preset_kwargs=dict(x=1, y=2, z=3, w=4)
+        ),
+        ([1, 2, 3], dict(z=3))
+    )
+    # Required preset argument is not provided.
+    with self.assertRaisesRegex(ValueError, 'Argument .* is not present.'):
+      callable_ext.PresetArgValue.resolve_args(
+          call_args=[],
+          call_kwargs=dict(),
+          positional_arg_names=['x'],
+          arg_defaults={
+              'x': callable_ext.PresetArgValue(),
+          },
+          preset_kwargs=dict()
+      )
+
+    # Include all preset kwargs.
+    self.assertEqual(
+        callable_ext.PresetArgValue.resolve_args(
+            call_args=[],
+            call_kwargs=dict(),
+            positional_arg_names=['x', 'y'],
+            arg_defaults={
+                'x': callable_ext.PresetArgValue(),
+                'y': 0,
+                'z': callable_ext.PresetArgValue(default=2),
+            },
+            preset_kwargs=dict(x=1, y=2, z=3, w=4),
+            include_all_preset_kwargs=True,
+        ),
+        ([1, 0], dict(z=3, w=4))
+    )
+
+  def test_preset_args(self):
+    @callable_ext.enable_preset_args()
+    def foo(
+        x=callable_ext.PresetArgValue(),
+        y=1,
+        *args,
+        z=callable_ext.PresetArgValue(default=2)
+    ):
+      del args
+      return x + y + z
+
+    with self.assertRaisesRegex(ValueError, 'Argument \'x\' is not present.'):
+      foo()
+
+    with callable_ext.preset_args(dict(x=1)):
+      self.assertEqual(foo(), 1 + 1 + 2)
+
+    # `y`` should not take precedence over the non-preset default.
+    with callable_ext.preset_args(dict(x=1, y=2)):
+      self.assertEqual(foo(), 1 + 1 + 2)
+
+    with callable_ext.preset_args(dict(x=1, y=2, z=3)):
+      self.assertEqual(foo(3), 3 + 1 + 3)
+
+    with callable_ext.preset_args(dict(x=1, y=2, z=3)):
+      self.assertEqual(foo(3, 3, z=4), 3 + 3 + 4)
+
+  def test_enable_preset_args(self):
+
+    # No side-effect if function does not have PresetArgValue.
+    def bar(x, y):
+      return x + y
+    self.assertIs(bar, callable_ext.enable_preset_args()(bar))
+
+    # `include_all_preset_kwargs` sets to False.
+    @callable_ext.enable_preset_args()
+    def baz(x, y=callable_ext.PresetArgValue(default=1), **kwargs):
+      return x + y + sum(kwargs.values())
+
+    with callable_ext.preset_args(dict(z=3, p=4)):
+      self.assertEqual(baz(1), 1 + 1)
+
+    # `include_all_prset_kwargs` is effective only when there is varkw.
+    @callable_ext.enable_preset_args(include_all_preset_kwargs=True)
+    def foo(x, y=callable_ext.PresetArgValue(default=1), **kwargs):
+      return x + y + sum(kwargs.values())
+
+    with callable_ext.preset_args(dict(z=3, p=4)):
+      self.assertEqual(foo(1), 1 + 1 + 3 + 4)
+
+    # `include_all_preset_kwargs` should be ignored if there is no varkw.
+    @callable_ext.enable_preset_args(include_all_preset_kwargs=True)
+    def fuz(x, y=callable_ext.PresetArgValue(default=1)):
+      return x + y
+
+    with callable_ext.preset_args(dict(y=2, z=3)):
+      self.assertEqual(fuz(1), 1 + 2)
+
+  def test_preset_args_nesting(self):
+    @callable_ext.enable_preset_args()
+    def foo(
+        x=callable_ext.PresetArgValue(),
+        y=1,
+        *,
+        z=callable_ext.PresetArgValue(default=2)
+    ):
+      return x + y + z
+
+    def bar(inherit_preset: bool = False, **kwargs):
+      with callable_ext.preset_args(
+          {k: v + 1 for k, v in kwargs.items()},
+          inherit_preset=inherit_preset
+      ):
+        return foo()
+
+    with callable_ext.preset_args(dict(x=1)):
+      self.assertEqual(foo(), 1 + 1 + 2)
+
+      self.assertEqual(bar(x=1), 2 + 1 + 2)
+      self.assertEqual(bar(x=1, z=2), 2 + 1 + 3)
+      self.assertEqual(bar(x=1, z=3), 2 + 1 + 4)
+
+      with self.assertRaisesRegex(ValueError, 'Argument \'x\' is not present.'):
+        bar()
+      self.assertEqual(bar(inherit_preset=True), 1 + 1 + 2)
+
+
 class CallWithOptionalKeywordArgsTest(unittest.TestCase):
   """Tests for typing.CallWithOptionalKeywordArgs."""