pyglove 0.4.5.dev202410020809__py3-none-any.whl → 0.4.5.dev202410100808__py3-none-any.whl

pyglove/core/typing/callable_ext.py

@@ -13,11 +13,250 @@
 # limitations under the License.
 """Callable extensions."""
 
-from typing import Any, Callable, List
+import contextlib
+import functools
+import inspect
+import types
+from typing import Any, Callable, Dict, Iterator, List, Optional, Sequence, Tuple, Union
 
+from pyglove.core import object_utils
 from pyglove.core.typing import callable_signature
 
 
+_TLS_KEY_PRESET_KWARGS = '__preset_kwargs__'
+
+
+class PresetArgValue(object_utils.Formattable):
+  """Value placeholder for arguments whose value will be provided by presets.
+
+  Example:
+
+    def foo(x, y=pg.PresetArgValue(default=1))
+      return x + y
+
+    with pg.preset_args(y=2):
+      print(foo(x=1))  # 3: y=2
+    print(foo(x=1))  # 2: y=1
+  """
+
+  def __init__(self, default: Any = object_utils.MISSING_VALUE):
+    self.default = default
+
+  @property
+  def has_default(self) -> bool:
+    return self.default != object_utils.MISSING_VALUE
+
+  def __eq__(self, other: Any) -> bool:
+    return isinstance(other, PresetArgValue) and (
+        self.default == other.default
+    )
+
+  def __ne__(self, other: Any) -> bool:
+    return not self.__eq__(other)
+
+  def format(self, *args, **kwargs):
+    return object_utils.kvlist_str(
+        [
+            ('default', self.default, object_utils.MISSING_VALUE),
+        ],
+        label='PresetArgValue',
+        *args,
+        **kwargs
+    )
+
+  @classmethod
+  def inspect(
+      cls, func: types.FunctionType) -> Dict[str, 'PresetArgValue']:
+    """Gets the PresetArgValue specified in a function's signature."""
+    assert inspect.isfunction(func), func
+    sig = inspect.signature(func)
+    preset_arg_markers = {}
+    for p in sig.parameters.values():
+      if isinstance(p.default, cls):
+        preset_arg_markers[p.name] = p.default
+    return preset_arg_markers
+
+  @classmethod
+  def resolve_args(
+      cls,
+      call_args: Tuple[Any, ...],
+      call_kwargs: Dict[str, Any],
+      positional_arg_names: Sequence[str],
+      arg_defaults: Dict[str, Any],
+      preset_kwargs: Dict[str, Any],
+      include_all_preset_kwargs: bool = False,
+  ) -> Tuple[Sequence[Any], Dict[str, Any]]:
+    """Resolves calling arguments passed to a method with presets."""
+    # Step 1: compute marked kwargs.
+    resolved_kwargs = {}
+    for arg_name, arg_default in arg_defaults.items():
+      if not isinstance(arg_default, PresetArgValue):
+        resolved_kwargs[arg_name] = arg_default
+        continue
+      if arg_name in preset_kwargs:
+        resolved_kwargs[arg_name] = preset_kwargs[arg_name]
+      elif arg_default.has_default:
+        resolved_kwargs[arg_name] = arg_default.default
+      else:
+        raise ValueError(
+            f'Argument {arg_name!r} is not present as a keyword argument '
+            'from the caller.'
+        )
+
+    # Step 2: add preset kwargs
+    if include_all_preset_kwargs:
+      for k, v in preset_kwargs.items():
+        if k not in resolved_kwargs:
+          resolved_kwargs[k] = v
+
+    # Step 3: merge call kwargs with resolved preset kwargs.
+    resolved_kwargs.update(call_kwargs)
+
+    # Step 3: remove resolved kwargs items as it's present in call args.
+    for i in range(len(call_args)):
+      if i >= len(positional_arg_names):
+        break
+      resolved_kwargs.pop(positional_arg_names[i], None)
+
+    # Step 4: convert kwargs back to postional arguments if applicable
+    resolved_args = call_args
+    if len(positional_arg_names) > len(call_args):
+      resolved_args = list(resolved_args)
+      for arg_name in positional_arg_names[len(call_args):]:
+        if arg_name not in resolved_kwargs:
+          break
+        arg_value = resolved_kwargs.pop(arg_name)
+        resolved_args.append(arg_value)
+    return resolved_args, resolved_kwargs
+
+
+class _ArgPresets:
+  """Preset argument collection."""
+
+  def __init__(self, presets: Optional[Dict[str, Dict[str, Any]]] = None):
+    self._presets: Dict[str, Dict[str, Any]] = presets or {}
+
+  def derive(
+      self,
+      kwargs: Dict[str, Any],
+      preset_name: str = 'global',
+      inherit_preset: Union[str, bool] = False
+  ) -> '_ArgPresets':
+    """Derives new presets from current presets."""
+    presets = self._presets.copy()   # Just do a shallow copy.
+    if isinstance(inherit_preset, bool) and inherit_preset:
+      inherit_preset = preset_name
+
+    if inherit_preset and inherit_preset in presets:
+      current_preset = presets[inherit_preset].copy()
+      current_preset.update(kwargs)
+    else:
+      current_preset = kwargs
+    presets[preset_name] = current_preset
+    return _ArgPresets(presets)
+
+  def get_preset(self, preset_name: str) -> Dict[str, Any]:
+    return self._presets.get(preset_name, {})
+
+
+@contextlib.contextmanager
+def preset_args(
+    kwargs: Dict[str, Any],
+    *,
+    preset_name: str = 'global',
+    inherit_preset: Union[str, bool] = False
+) -> Iterator[Dict[str, Any]]:
+  """Context manager to enable calling with user kwargs.
+
+  Args:
+    kwargs: The preset kwargs to be used by preset-enabled functions within
+      the context.
+    preset_name: The name of the preset to specify kwargs.
+      `enable_preset_args` allows users to pass a preset name, which will be
+      used to identify the present to be used.
+    inherit_preset: The name of the preset defined by the parent context to
+      inherit kwargs from. Or a boolean to indicate whether to inherit a
+      parent preset of the same name.
+
+  Yields:
+    Current preset kwargs.
+  """
+
+  parent_presets = object_utils.thread_local_peek(
+      _TLS_KEY_PRESET_KWARGS, _ArgPresets()
+  )
+  current_preset = parent_presets.derive(kwargs, preset_name, inherit_preset)
+  object_utils.thread_local_push(_TLS_KEY_PRESET_KWARGS, current_preset)
+  try:
+    yield current_preset
+  finally:
+    object_utils.thread_local_pop(_TLS_KEY_PRESET_KWARGS, None)
+
+
+def enable_preset_args(
+    include_all_preset_kwargs: bool = False,
+    preset_name: str = 'global'
+) -> Callable[[types.FunctionType], types.FunctionType]:
+  """Decorator for functions that maybe use preset argument values.
+
+  Usage::
+
+    @pg.typing.enable_preset_args
+    def foo(x, y=pg.typing.PresetArgValue(default=1)):
+      return x + y
+
+    with pg.typing.preset_args(y=2):
+      print(foo(x=1))  # 3: y=2
+    print(foo(x=1))  # 2: y=1
+
+  Args:
+    include_all_preset_kwargs: Whether to include all preset kwargs (even
+      not makred as `PresetArgValue`) when callng the function.
+    preset_name: The name of the preset to specify kwargs.
+
+  Returns:
+    A decorated function that could consume the preset argument values.
+  """
+  def decorator(func):
+    sig = inspect.signature(func)
+    positional_arg_names = [
+        p.name for p in sig.parameters.values()
+        if p.kind == inspect.Parameter.POSITIONAL_OR_KEYWORD
+    ]
+
+    arg_defaults = {}
+    has_preset_value = False
+    has_varkw = False
+    for p in sig.parameters.values():
+      if p.kind == inspect.Parameter.VAR_KEYWORD:
+        has_varkw = True
+        continue
+      if p.kind == inspect.Parameter.VAR_POSITIONAL:
+        continue
+      if p.default == inspect.Parameter.empty:
+        continue
+      if isinstance(p.default, PresetArgValue):
+        has_preset_value = True
+      arg_defaults[p.name] = p.default
+
+    if has_preset_value:
+      @functools.wraps(func)
+      def _func(*args, **kwargs):
+        # Map positional arguments to keyword arguments.
+        presets = object_utils.thread_local_peek(
+            _TLS_KEY_PRESET_KWARGS, None
+        )
+        preset_kwargs = presets.get_preset(preset_name) if presets else {}
+        args, kwargs = PresetArgValue.resolve_args(
+            args, kwargs, positional_arg_names, arg_defaults, preset_kwargs,
+            include_all_preset_kwargs=include_all_preset_kwargs and has_varkw
+        )
+        return func(*args, **kwargs)
+      return _func
+    return func
+  return decorator
+
+
 class CallableWithOptionalKeywordArgs:
   """Helper class for invoking callable objects with optional keyword args.
 

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."""
 

pyglove/core/typing/inspect.py

@@ -14,6 +14,7 @@
 """Utility module for inspecting generics types."""
 
 import inspect
+import sys
 import typing
 from typing import Any, Callable, Optional, Tuple, Type, Union
 
@@ -116,6 +117,68 @@ def get_type_args(
     return ()
 
 
+def get_outer_class(
+    cls: Type[Any],
+    base_cls: Union[Type[Any], Tuple[Type[Any], ...], None] = None,
+    immediate: bool = False,
+) -> Optional[Type[Any]]:
+  """Returns the outer class.
+
+  Example::
+
+    class A:
+      pass
+
+    class A1:
+      class B:
+        class C:
+          ...
+
+    pg.typing.outer_class(B) is A1
+    pg.typing.outer_class(C) is B
+    pg.typing.outer_class(C, base_cls=A) is None
+    pg.typing.outer_class(C, base_cls=A1) is None
+
+  Args:
+    cls: The class to get the outer class for.
+    base_cls: The base class of the outer class. If provided, an outer class
+      that is not a subclass of `base_cls` will be returned as None.
+    immediate: Whether to return the immediate outer class or a class in the
+      nesting hierarchy that is a subclass of `base_cls`. Applicable when
+      `base_cls` is not None.
+
+  Returns:
+    The outer class of `cls`. None if cannot find one or the outer class is
+      not a subclass of `base_cls`.
+  """
+  if '<locals>' in cls.__qualname__:
+    raise ValueError(
+        'Cannot find the outer class for locally defined class '
+        f'{cls.__qualname__!r}'
+    )
+
+  names = cls.__qualname__.split('.')
+  if len(names) < 2:
+    return None
+
+  parent = sys.modules[cls.__module__]
+  symbols = []
+  for name in names[:-1]:
+    symbol = getattr(parent, name, None)
+    if symbol is None:
+      return None
+    assert inspect.isclass(symbol), symbol
+    symbols.append(symbol)
+    parent = symbol
+
+  for symbol in reversed(symbols):
+    if immediate:
+      return symbol if not base_cls or issubclass(symbol, base_cls) else None
+    if not base_cls or issubclass(symbol, base_cls):
+      return symbol
+  return None
+
+
 def callable_eq(
     x: Optional[Callable[..., Any]], y: Optional[Callable[..., Any]]
 ) -> bool:

pyglove/core/typing/inspect_test.py

@@ -16,8 +16,10 @@
 from typing import Any, Generic, TypeVar
 import unittest
 
+from pyglove.core.typing import callable_signature
 from pyglove.core.typing import inspect
 
+
 XType = TypeVar('XType')
 YType = TypeVar('YType')
 
@@ -50,6 +52,16 @@ class D(C):
   pass
 
 
+class AA:
+  pass
+
+
+class AA1(AA):
+  class BB1:
+    class CC1:
+      pass
+
+
 class InspectTest(unittest.TestCase):
 
   def test_issubclass(self):
@@ -141,6 +153,33 @@ class InspectTest(unittest.TestCase):
     self.assertEqual(inspect.get_type_args(C, A), (str, int))
     self.assertEqual(inspect.get_type_args(C, B), (Str,))
 
+  def test_outer_class(self):
+    class Foo:
+      pass
+
+    with self.assertRaisesRegex(ValueError, '.* locally defined class'):
+      inspect.get_outer_class(Foo)
+
+    self.assertIsNone(inspect.get_outer_class(AA))
+    self.assertIs(inspect.get_outer_class(AA1.BB1), AA1)
+    self.assertIs(inspect.get_outer_class(AA1.BB1, AA), AA1)
+    self.assertIs(inspect.get_outer_class(AA1.BB1, A), None)
+    self.assertIs(inspect.get_outer_class(AA1.BB1.CC1), AA1.BB1)
+    self.assertIsNone(
+        inspect.get_outer_class(AA1.BB1.CC1, base_cls=AA, immediate=True)
+    )
+    self.assertIs(inspect.get_outer_class(AA1.BB1.CC1, AA), AA1)
+    self.assertIs(
+        inspect.get_outer_class(callable_signature.Argument.Kind),
+        callable_signature.Argument
+    )
+
+    class Bar:
+      pass
+
+    Bar.__qualname__ = 'NonExist.Bar'
+    self.assertIsNone(inspect.get_outer_class(Bar))
+
   def test_callable_eq(self):
     def foo(unused_x):
       pass

pyglove/core/views/__init__.py

@@ -0,0 +1,30 @@
+# Copyright 2024 The PyGlove Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+"""PyGlove views."""
+
+from pyglove.core.views import base
+from pyglove.core.views import html
+
+View = base.View
+view = base.view
+
+# Pytype annotation.
+NodeFilter = base.NodeFilter
+
+Html = html.Html
+HtmlView = html.HtmlView
+HtmlTreeView = html.HtmlTreeView
+
+to_html = html.to_html
+to_html_str = html.to_html_str