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

pyglove/core/typing/annotation_conversion.py

@@ -18,7 +18,7 @@ import inspect
 import types
 import typing
 
-from pyglove.core import object_utils
+from pyglove.core import utils
 from pyglove.core.typing import annotated
 from pyglove.core.typing import class_schema
 from pyglove.core.typing import inspect as pg_inspect
@@ -40,7 +40,8 @@ def _field_from_annotation(
     annotation: typing.Any,
     description: typing.Optional[str] = None,
     metadata: typing.Optional[typing.Dict[str, typing.Any]] = None,
-    auto_typing=True,
+    auto_typing: bool = True,
+    parent_module: typing.Optional[types.ModuleType] = None
 ) -> class_schema.Field:
   """Creates a field from Python annotation."""
   if isinstance(annotation, annotated.Annotated):
@@ -55,7 +56,9 @@ def _field_from_annotation(
   return class_schema.create_field(
       field_spec,
       auto_typing=auto_typing,
-      accept_value_as_annotation=False)
+      accept_value_as_annotation=False,
+      parent_module=parent_module
+  )
 
 
 def _value_spec_from_default_value(
@@ -70,7 +73,7 @@ def _value_spec_from_default_value(
   elif isinstance(value, tuple):
     value_spec = vs.Tuple(
         [_value_spec_from_default_value(elem, False) for elem in value])
-  elif inspect.isfunction(value) or isinstance(value, object_utils.Functor):
+  elif inspect.isfunction(value) or isinstance(value, utils.Functor):
     value_spec = vs.Callable()
   elif not isinstance(value, type):
     value_spec = vs.Object(type(value))
@@ -84,7 +87,9 @@ def _value_spec_from_default_value(
 
 def _value_spec_from_type_annotation(
     annotation: typing.Any,
-    accept_value_as_annotation: bool) -> class_schema.ValueSpec:
+    accept_value_as_annotation: bool,
+    parent_module: typing.Optional[types.ModuleType] = None
+) -> class_schema.ValueSpec:
   """Creates a value spec from type annotation."""
   if annotation is bool:
     return vs.Bool()
@@ -100,11 +105,15 @@ def _value_spec_from_type_annotation(
   origin = typing.get_origin(annotation) or annotation
   args = list(typing.get_args(annotation))
 
+  def _sub_value_spec_from_annotation(
+      annotation: typing.Any) -> class_schema.ValueSpec:
+    return _value_spec_from_type_annotation(
+        annotation, accept_value_as_annotation, parent_module)
+
   # Handling list.
   if origin in (list, typing.List):
     return vs.List(
-        _value_spec_from_annotation(
-            args[0], True)) if args else vs.List(vs.Any())
+        _sub_value_spec_from_annotation(args[0])) if args else vs.List(vs.Any())
   # Handling tuple.
   elif origin in (tuple, typing.Tuple):
     if not args:
@@ -115,16 +124,15 @@ def _value_spec_from_type_annotation(
           raise TypeError(
               f'Tuple with ellipsis should have exact 2 type arguments. '
               f'Encountered: {annotation}.')
-        return vs.Tuple(_value_spec_from_type_annotation(args[0], False))
-      return vs.Tuple([_value_spec_from_type_annotation(arg, False)
-                       for arg in args])
+        return vs.Tuple(_sub_value_spec_from_annotation(args[0]))
+      return vs.Tuple([_sub_value_spec_from_annotation(arg) for arg in args])
   # Handling sequence.
   elif origin in (collections.abc.Sequence,):
-    elem = _value_spec_from_annotation(args[0], True) if args else vs.Any()
+    elem = _sub_value_spec_from_annotation(args[0]) if args else vs.Any()
     return vs.Union([vs.List(elem), vs.Tuple(elem)])
   # Handling literals.
   elif origin is typing.Literal:
-    return vs.Enum(object_utils.MISSING_VALUE, args)
+    return vs.Enum(utils.MISSING_VALUE, args)
   # Handling dict.
   elif origin in (dict, typing.Dict, collections.abc.Mapping):
     if not args:
@@ -133,8 +141,7 @@ def _value_spec_from_type_annotation(
     if args[0] not in (str, typing.Text):
       raise TypeError(
           'Dict type field with non-string key is not supported.')
-    elem_value_spec = _value_spec_from_type_annotation(
-        args[1], accept_value_as_annotation=False)
+    elem_value_spec = _sub_value_spec_from_annotation(args[1])
     return vs.Dict([(ks.StrKey(), elem_value_spec)])
   elif origin is collections.abc.Callable:
     arg_specs = []
@@ -149,13 +156,8 @@ def _value_spec_from_type_annotation(
       #   Callable[int, Any] => Callable[[int], Any]
       #   Callable[(int, int), Any] => Callable[[int, int], Any]
       if isinstance(args[0], list):
-        arg_specs = [
-            _value_spec_from_type_annotation(
-                arg, accept_value_as_annotation=False)
-            for arg in args[0]
-        ]
-      return_spec = _value_spec_from_type_annotation(
-          args[1], accept_value_as_annotation=False)
+        arg_specs = [_sub_value_spec_from_annotation(arg) for arg in args[0]]
+      return_spec = _sub_value_spec_from_annotation(args[1])
     return vs.Callable(arg_specs, returns=return_spec)
   # Handling type
   elif origin is type or (annotation in (typing.Type, type)):
@@ -169,20 +171,32 @@ def _value_spec_from_type_annotation(
     if optional:
       args.remove(_NoneType)
     if len(args) == 1:
-      spec = _value_spec_from_annotation(args[0], True)
+      spec = _sub_value_spec_from_annotation(args[0])
     else:
-      spec = vs.Union([_value_spec_from_annotation(x, True) for x in set(args)])
+      spec = vs.Union([_sub_value_spec_from_annotation(x) for x in args])
     if optional:
       spec = spec.noneable()
     return spec
+  elif origin is typing.Final:
+    return _value_spec_from_type_annotation(
+        args[0],
+        accept_value_as_annotation=False
+    ).freeze(vs._FROZEN_VALUE_PLACEHOLDER)  # pylint: disable=protected-access
   elif isinstance(annotation, typing.ForwardRef):
-    return vs.Object(annotation.__forward_arg__)
+    annotation = annotation.__forward_arg__
+    if parent_module is not None:
+      annotation = class_schema.ForwardRef(parent_module, annotation)
+    return vs.Object(annotation)
+  elif isinstance(annotation, class_schema.ForwardRef):
+    return vs.Object(annotation)
   # Handling class.
   elif (
       inspect.isclass(annotation)
       or pg_inspect.is_generic(annotation)
       or (isinstance(annotation, str) and not accept_value_as_annotation)
   ):
+    if isinstance(annotation, str) and parent_module is not None:
+      annotation = class_schema.ForwardRef(parent_module, annotation)
     return vs.Object(annotation)
 
   if accept_value_as_annotation:
@@ -204,8 +218,9 @@ def _any_spec_with_annotation(annotation: typing.Any) -> vs.Any:
 
 def _value_spec_from_annotation(
     annotation: typing.Any,
-    auto_typing=False,
-    accept_value_as_annotation=False
+    auto_typing: bool = False,
+    accept_value_as_annotation: bool = False,
+    parent_module: typing.Optional[types.ModuleType] = None
     ) -> class_schema.ValueSpec:
   """Creates a value spec from annotation."""
   if isinstance(annotation, class_schema.ValueSpec):
@@ -220,7 +235,8 @@ def _value_spec_from_annotation(
 
   if auto_typing:
     return _value_spec_from_type_annotation(
-        annotation, accept_value_as_annotation)
+        annotation, accept_value_as_annotation, parent_module
+    )
   else:
     value_spec = None
     if accept_value_as_annotation:

pyglove/core/typing/annotation_conversion_test.py

@@ -14,6 +14,7 @@
 """Tests for pyglove.core.typing.annotation_conversion."""
 
 import inspect
+import sys
 import typing
 import unittest
 
@@ -22,6 +23,7 @@ from pyglove.core.typing import annotation_conversion
 from pyglove.core.typing import key_specs as ks
 from pyglove.core.typing import value_specs as vs
 from pyglove.core.typing.class_schema import Field
+from pyglove.core.typing.class_schema import ForwardRef
 from pyglove.core.typing.class_schema import ValueSpec
 
 
@@ -74,6 +76,14 @@ class FieldFromAnnotationTest(unittest.TestCase):
         Field.from_annotation('x', str, 'A str', dict(x=1), auto_typing=True),
         Field('x', vs.Str(), 'A str', dict(x=1)))
 
+  def test_with_parent_module(self):
+    self.assertEqual(
+        Field.from_annotation(
+            'x', 'ValueSpecBase', auto_typing=True, parent_module=vs
+        ),
+        Field('x', vs.Object(vs.ValueSpecBase))
+    )
+
 
 class ValueSpecFromAnnotationTest(unittest.TestCase):
   """Tests for ValueSpec.fromAnnotation."""
@@ -248,6 +258,11 @@ class ValueSpecFromAnnotationTest(unittest.TestCase):
     self.assertEqual(
         ValueSpec.from_annotation(
             typing.ForwardRef('Foo'), True), vs.Object(Foo))
+    self.assertEqual(
+        ValueSpec.from_annotation(
+            ForwardRef(sys.modules[__name__], 'Foo'),
+            True
+        ), vs.Object(Foo))
 
   def test_generic_class(self):
     X = typing.TypeVar('X')
@@ -291,6 +306,14 @@ class ValueSpecFromAnnotationTest(unittest.TestCase):
           ValueSpec.from_annotation(int | str, True),
           vs.Union([vs.Int(), vs.Str()]))
 
+  def test_final(self):
+    self.assertEqual(
+        ValueSpec.from_annotation(
+            typing.Final[int], True
+        ).set_default(1),
+        vs.Int().freeze(1)
+    )
+
 
 if __name__ == '__main__':
   unittest.main()

pyglove/core/typing/callable_ext.py

@@ -13,11 +13,248 @@
 # 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 utils
 from pyglove.core.typing import callable_signature
 
 
+_TLS_KEY_PRESET_KWARGS = '__preset_kwargs__'
+
+
+class PresetArgValue(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 = utils.MISSING_VALUE):
+    self.default = default
+
+  @property
+  def has_default(self) -> bool:
+    return self.default != 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 utils.kvlist_str(
+        [
+            ('default', self.default, 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 = utils.thread_local_peek(
+      _TLS_KEY_PRESET_KWARGS, _ArgPresets()
+  )
+  current_preset = parent_presets.derive(kwargs, preset_name, inherit_preset)
+  utils.thread_local_push(_TLS_KEY_PRESET_KWARGS, current_preset)
+  try:
+    yield current_preset
+  finally:
+    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 = 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.
 
@@ -31,8 +268,9 @@ class CallableWithOptionalKeywordArgs:
   def __init__(self,
                func: Callable[..., Any],
                optional_keywords: List[str]):
-    sig = callable_signature.get_signature(func)
-    absent_keywords = None
+    sig = callable_signature.signature(
+        func, auto_typing=False, auto_doc=False
+    )
 
     # Check for variable keyword arguments.
     if sig.has_varkw: