Add additional support for representing objects in YAML serialization.

fiddle/experimental/testdata/yaml_serialization_diamond.yaml

@@ -1,14 +1,10 @@
 !fdl.Config
-__fn_or_cls__:
-  module: __main__
-  name: Foo
+__fn_or_cls__: !class '__main__.Foo'
 a: 2
 b: root
 c:
 - &id001 !fdl.Config
-  __fn_or_cls__:
-    module: __main__
-    name: Foo
+  __fn_or_cls__: !class '__main__.Foo'
   a: 1
   b: shared
   c: null

fiddle/experimental/yaml_serialization.py

@@ -50,10 +50,7 @@ def _config_representer(dumper, data, type_name="fdl.Config"):
     raise ValueError("It is not supported to dump objects of functions/classes "
                      "that have a __fn_or_cls__ parameter.")
 
-  value["__fn_or_cls__"] = {
-      "module": inspect.getmodule(data.__fn_or_cls__).__name__,
-      "name": data.__fn_or_cls__.__qualname__,
-  }
+  value["__fn_or_cls__"] = data.__fn_or_cls__
   return dumper.represent_mapping(f"!{type_name}", value)
 
 
@@ -67,26 +64,44 @@ def _fixture_representer(dumper, data):
 
 
 def _taggedvalue_representer(dumper, data):
-  return dumper.represent_mapping("!fdl.TaggedValue", {
+  return dumper.represent_mapping("!fdl.TaggedValueCls", {
       "tags": [tag.name for tag in data.tags],
       "value": data.value,
   })
 
 
+def _custom_object_representer(dumper, data):
+  """Representer for arbitrary Python objects."""
+  if inspect.isclass(data) or inspect.isfunction(data):
+    return dumper.represent_scalar(
+        "!class" if inspect.isclass(data) else "!function",
+        f"{data.__module__}.{data.__qualname__}",
+        style="'")
+  else:
+    attrs = {k: v for k, v in data.__dict__.items() if not k.startswith("_")}
+    attrs["__type__"] = f"{type(data).__module__}.{type(data).__qualname__}"
+    return dumper.represent_mapping("!object", attrs)
+
+
 yaml.SafeDumper.add_representer(config.Config, _config_representer)
 yaml.SafeDumper.add_representer(config.Partial, _partial_representer)
 yaml.SafeDumper.add_representer(fixture.Fixture, _fixture_representer)
 yaml.SafeDumper.add_representer(tagging.TaggedValueCls,
                                 _taggedvalue_representer)
+yaml.SafeDumper.add_representer(None, _custom_object_representer)
+
+
+class _CustomSafeDumper(yaml.SafeDumper):
+
+  def ignore_aliases(self, data):
+    return (super().ignore_aliases(data) or inspect.isclass(data) or
+            inspect.isfunction(data))
 
 
 def dump_yaml(value: Any) -> str:
   """Returns the YAML serialization of `value`.
 
   Args:
     value: The value to serialize.
-
-  Raises:
-    PyrefError: If an error is encountered while serializing a Python reference.
   """
-  return yaml.safe_dump(value)
+  return yaml.dump(value, Dumper=_CustomSafeDumper)

fiddle/experimental/yaml_serialization_test.py

@@ -16,6 +16,7 @@
 """Tests for yaml_serialization."""
 import dataclasses
 import pathlib
+import textwrap
 from typing import Any
 
 from absl.testing import absltest
@@ -34,12 +35,7 @@ def _testdata_dir():
 
 def _config_constructor(loader, node):
   arguments = loader.construct_mapping(node, deep=True)
-  fn_or_cls_pyref = arguments.pop("__fn_or_cls__")
-  fn_or_cls = serialization.import_symbol(
-      serialization.DefaultPyrefPolicy(),
-      fn_or_cls_pyref["module"],
-      fn_or_cls_pyref["name"],
-  )
+  fn_or_cls = arguments.pop("__fn_or_cls__")
   return fdl.Config(fn_or_cls, **arguments)
 
 
@@ -51,6 +47,13 @@ def _fixture_constructor(loader, node):
   return fixture.Fixture(_config_constructor(loader, node))
 
 
+def _fn_or_cls_constructor(loader, node):
+  del loader
+  module, name = node.value.rsplit(".", 1)
+  policy = serialization.DefaultPyrefPolicy()
+  return serialization.import_symbol(policy, module, name)
+
+
 class SemiSafeLoader(yaml.SafeLoader):
   """Intermediate class that can load Fiddle configs."""
 
@@ -73,6 +76,8 @@ def load_yaml_test_only(serialized: str) -> Any:
   SemiSafeLoader.add_constructor("!fdl.Partial", _partial_constructor)
   SemiSafeLoader.add_constructor("!fiddle.experimental.Fixture",
                                  _fixture_constructor)
+  SemiSafeLoader.add_constructor("!function", _fn_or_cls_constructor)
+  SemiSafeLoader.add_constructor("!class", _fn_or_cls_constructor)
   return yaml.load(serialized, Loader=SemiSafeLoader)
 
 
@@ -136,12 +141,12 @@ def test_dump_fixture(self):
     self.assertEqual(loaded, config)
 
   def test_dump_tagged_value(self):
-    self.assertRegex(
-        yaml_serialization.dump_yaml(value=FakeTag.new(2)).strip(),
-        r"""!fdl\.TaggedValue
-tags:
-- [\w\d_\.]+\.FakeTag
-value: 2""")
+    regex = textwrap.dedent(r"""
+        !fdl\.TaggedValueCls
+        tags:
+        - [\w\d_\.]+\.FakeTag
+        value: 2""").lstrip()
+    self.assertRegex(yaml_serialization.dump_yaml(value=FakeTag.new(2)), regex)
 
   def test_dump_diamond(self):
     shared = fdl.Config(Foo, a=1, b="shared", c=None)
@@ -155,6 +160,16 @@ def test_dump_diamond(self):
     loaded = load_yaml_test_only(serialized)
     self.assertDagEqual(config, loaded)
 
+  def test_dump_custom_object(self):
+    serialized = yaml_serialization.dump_yaml(value=Foo(1, "a", None))
+    expected = textwrap.dedent("""
+        !object
+        __type__: __main__.Foo
+        a: 1
+        b: a
+        c: null""")
+    self.assertEqual(expected.strip(), serialized.strip())
+
 
 if __name__ == "__main__":
   absltest.main()