From 9416caade49af4b8f2bfabe4186267160e4a4fc1 Mon Sep 17 00:00:00 2001
From: Darsey Litzenberger <dlitz@dlitz.net>
Date: Fri, 2 Aug 2024 01:03:02 -0600
Subject: [PATCH] feat(python): add pretty-printing of protobuf messages in
 IPython

---
 python/.changelog.d/4076.added   |  1 +
 python/src/trezorlib/protobuf.py | 36 ++++++++++++++++++++++++++++++++
 2 files changed, 37 insertions(+)
 create mode 100644 python/.changelog.d/4076.added

diff --git a/python/.changelog.d/4076.added b/python/.changelog.d/4076.added
new file mode 100644
index 00000000000..faeb0765188
--- /dev/null
+++ b/python/.changelog.d/4076.added
@@ -0,0 +1 @@
+Added pretty-printing of protobuf messages in IPython (`_repr_pretty_`)
diff --git a/python/src/trezorlib/protobuf.py b/python/src/trezorlib/protobuf.py
index 836716a5422..9109990e452 100644
--- a/python/src/trezorlib/protobuf.py
+++ b/python/src/trezorlib/protobuf.py
@@ -35,6 +35,9 @@
 
 import typing_extensions as tx
 
+if t.TYPE_CHECKING:
+    from IPython.lib.pretty import RepresentationPrinter  # noqa: I900
+
 T = t.TypeVar("T", bound=type)
 MT = t.TypeVar("MT", bound="MessageType")
 
@@ -271,6 +274,39 @@ def __repr__(self) -> str:
             d[key] = value
         return f"<{self.__class__.__name__}: {d}>"
 
+    def _repr_pretty_(self, p: RepresentationPrinter, cycle: bool) -> None:
+        """prettier version of __repr__ for IPython
+
+        This pretty-prints/indents the object when displayed in IPython,
+        for example:
+
+            <PrevInput: {'prev_hash': b'xyzasdf',
+                         'prev_index': 1,
+                         'script_sig': b'abcdef',
+                         'sequence': 42,
+                         'decred_tree': 21}>
+
+        The API is for this method is described in the IPython docs:
+        https://ipython.readthedocs.io/en/8.26.0/api/generated/IPython.lib.pretty.html
+        """
+        prefix = f"<{self.__class__.__name__}: {{"
+        if cycle:
+            p.text(f"{prefix} ...>")
+            return
+        with p.group(len(prefix), prefix, f"}}>"):  # noqa: F541
+            itemsiter = (
+                (key, value)
+                for key, value in self.__dict__.items()
+                if not (value is None or value == [])
+            )
+            for i, (key, value) in enumerate(itemsiter):
+                if i:
+                    p.text(",")
+                    p.breakable()
+                subprefix = f"{key!r}: "
+                with p.group(len(subprefix), subprefix, ""):
+                    p.pretty(value)
+
     def ByteSize(self) -> int:
         data = BytesIO()
         dump_message(data, self)