From fd6ea8dd035d1dc07164c58f0dabe1ad8c063c05 Mon Sep 17 00:00:00 2001
From: Daiyi Peng <daiyip@google.com>
Date: Thu, 4 Jan 2024 05:46:24 -0800
Subject: [PATCH] Support markdown quotes on the repr/str of Formattable
 objects.

This allows users to use `pg.str_format` and `pg.repr_format` context manager to add markdown quotes to the repr/str of symbolic objects, which serves the scenario when embedding objects within natural language string.

PiperOrigin-RevId: 595679721
---
 pyglove/core/__init__.py                      |  3 ++
 pyglove/core/object_utils/common_traits.py    | 22 +++++++++++--
 .../core/object_utils/common_traits_test.py   | 31 ++++++++++++++++---
 3 files changed, 49 insertions(+), 7 deletions(-)

diff --git a/pyglove/core/__init__.py b/pyglove/core/__init__.py
index 61c58c5..00a9381 100644
--- a/pyglove/core/__init__.py
+++ b/pyglove/core/__init__.py
@@ -279,6 +279,9 @@
 MISSING_VALUE = object_utils.MISSING_VALUE
 
 Formattable = object_utils.Formattable
+repr_format = object_utils.repr_format
+str_format = object_utils.str_format
+
 MaybePartial = object_utils.MaybePartial
 JSONConvertible = object_utils.JSONConvertible
 DocStr = object_utils.DocStr
diff --git a/pyglove/core/object_utils/common_traits.py b/pyglove/core/object_utils/common_traits.py
index ed423ab..2aee2bd 100644
--- a/pyglove/core/object_utils/common_traits.py
+++ b/pyglove/core/object_utils/common_traits.py
@@ -76,13 +76,31 @@ def __str__(self) -> str:
     """Returns the full (maybe multi-line) representation of this object."""
     kwargs = dict(self.__str_format_kwargs__)
     kwargs.update(thread_local.thread_local_kwargs(_TLS_STR_FORMAT_KWARGS))
-    return self.format(**kwargs)
+    return self._maybe_quote(self.format(**kwargs), **kwargs)
 
   def __repr__(self) -> str:
     """Returns a single-line representation of this object."""
     kwargs = dict(self.__repr_format_kwargs__)
     kwargs.update(thread_local.thread_local_kwargs(_TLS_REPR_FORMAT_KWARGS))
-    return self.format(**kwargs)
+    return self._maybe_quote(self.format(**kwargs), **kwargs)
+
+  def _maybe_quote(
+      self,
+      s: str,
+      *,
+      compact: bool = False,
+      root_indent: int = 0,
+      markdown: bool = False,
+      **kwargs
+  ) -> str:
+    """Maybe quote the formatted string with markdown."""
+    del kwargs
+    if not markdown or root_indent > 0:
+      return s
+    if compact:
+      return f'`{s}`'
+    else:
+      return f'\n```\n{s}\n```\n'
 
 
 class MaybePartial(metaclass=abc.ABCMeta):
diff --git a/pyglove/core/object_utils/common_traits_test.py b/pyglove/core/object_utils/common_traits_test.py
index dfc8d18..0595d84 100644
--- a/pyglove/core/object_utils/common_traits_test.py
+++ b/pyglove/core/object_utils/common_traits_test.py
@@ -19,10 +19,24 @@
 
 class Foo(common_traits.Formattable):
 
-  def format(self, compact: bool = False, verbose: bool = True, **kwargs):
+  def format(
+      self, compact: bool = False, verbose: bool = True, **kwargs):
     return f'{self.__class__.__name__}(compact={compact}, verbose={verbose})'
 
 
+class Bar(common_traits.Formattable):
+
+  def __init__(self, foo: Foo):
+    self._foo = foo
+
+  def format(
+      self, compact: bool = False, verbose: bool = True,
+      root_indent: int = 0, **kwargs):
+    foo_str = self._foo.format(
+        compact=compact, verbose=verbose, root_indent=root_indent + 1)
+    return f'{self.__class__.__name__}(foo={foo_str})'
+
+
 class FormattableTest(unittest.TestCase):
 
   def test_formattable(self):
@@ -31,13 +45,13 @@ def test_formattable(self):
     self.assertEqual(str(foo), 'Foo(compact=False, verbose=True)')
 
   def test_formattable_with_custom_format(self):
-    class Bar(Foo):
+    class Baz(Foo):
       __str_format_kwargs__ = {'compact': False, 'verbose': False}
       __repr_format_kwargs__ = {'compact': True, 'verbose': False}
 
-    bar = Bar()
-    self.assertEqual(repr(bar), 'Bar(compact=True, verbose=False)')
-    self.assertEqual(str(bar), 'Bar(compact=False, verbose=False)')
+    bar = Baz()
+    self.assertEqual(repr(bar), 'Baz(compact=True, verbose=False)')
+    self.assertEqual(str(bar), 'Baz(compact=False, verbose=False)')
 
   def test_formattable_with_context_managers(self):
     foo = Foo()
@@ -46,6 +60,13 @@ def test_formattable_with_context_managers(self):
         self.assertEqual(repr(foo), 'Foo(compact=False, verbose=True)')
         self.assertEqual(str(foo), 'Foo(compact=False, verbose=False)')
 
+    bar = Bar(foo)
+    with common_traits.repr_format(markdown=True):
+      self.assertEqual(repr(bar), '`Bar(foo=Foo(compact=True, verbose=True))`')
+    with common_traits.str_format(markdown=True):
+      self.assertEqual(
+          str(bar), '\n```\nBar(foo=Foo(compact=False, verbose=True))\n```\n')
+
 
 class ExplicitlyOverrideTest(unittest.TestCase):