Improve contractions.

Author: MarkDaoust

google/generativeai/string_utils.py

@@ -16,6 +16,7 @@
 
 import dataclasses
 import pprint
+import re
 import reprlib
 import textwrap
 
@@ -33,17 +34,33 @@ def prettyprint(cls):
     return cls
 
 
+repr = reprlib.Repr()
+
+
 @reprlib.recursive_repr()
 def _prettyprint(self):
-    """You can't use `__str__ = pprint.pformat`. That causes a recursion error.
+    """A dataclass prettyprint functrion you can use in __str__or __repr__.
+
+    Note: You can't set `__str__ = pprint.pformat`. That causes a recursion error.
+
+    Mostly identical to pprint but:
 
-    This works, but it doesn't handle objects that reference themselves.
+    * This will contract long lists and dicts to [...] and {...}.
+    * This will contract object reprs ClassName(...).
     """
     fields = []
     for f in dataclasses.fields(self):
         s = pprint.pformat(getattr(self, f.name))
+        class_re = r"^(\w+)\(.*\)$"
         if s.count("\n") >= 10:
-            s = "..."
+            if s.startswith("["):
+                s = "[...]"
+            elif s.startswith("{"):
+                s = "{...}"
+            elif re.match(class_re, s, flags=re.DOTALL):
+                s = re.sub(class_re, r"\1(...)", s, flags=re.DOTALL)
+            else:
+                s = "..."
         else:
             width = len(f.name) + 1
             s = textwrap.indent(s, " " * width).lstrip(" ")

tests/test_string_utils.py

@@ -50,13 +50,13 @@ def test_simple(self):
         self.assertEqual(pprint.pformat(m), result)
         self.assertEqual(repr(m), result)
 
-    def test_long(self):
+    def test_long_list(self):
         m = MyClass(a=1, b=1 / 3, c=[1, 2, 3] * 10, d={"a": 1, "b": 2})
         expected = textwrap.dedent(
             """
             MyClass(a=1,
                     b=0.3333333333333333,
-                    c=...,
+                    c=[...],
                     d={'a': 1, 'b': 2})"""
         )[1:]
         self.assertEqual(expected, str(m))
@@ -80,6 +80,22 @@ def test_nested(self):
         self.assertEqual(pprint.pformat(m2), result)
         self.assertEqual(repr(m2), result)
 
+    def test_long_obj(self):
+        m = MyClass(a=1, b=1 / 3, c=[1, 2, 3], d=None)
+        m = MyClass(a=1, b=1 / 3, c=[1, 2, 3], d=m)
+        m = MyClass(a=1, b=1 / 3, c=[1, 2, 3], d=m)
+        m = MyClass(a=1, b=1 / 3, c=[1, 2, 3], d=m)
+        m = MyClass(a=1, b=1 / 3, c=[1, 2, 3], d=m)
+
+        expected = textwrap.dedent(
+            """
+            MyClass(a=1,
+                    b=0.3333333333333333,
+                    c=[1, 2, 3],
+                    d=MyClass(...))"""
+        )[1:]
+        self.assertEqual(expected, str(m))
+
     def test_recursive(self):
         m = MyClass(a=1, b=1 / 3, c=[1, 2, 3], d=None)
         m.d = m