From ca8ebacab9a707182f5160453202c35b42aa32be Mon Sep 17 00:00:00 2001
From: matejcik <ja@matejcik.cz>
Date: Wed, 9 May 2018 18:09:31 +0200
Subject: [PATCH] logging: convenience function to set up protobuf logging

---
 trezorctl        |  8 ++------
 trezorlib/log.py | 23 +++++++++++++++++++----
 2 files changed, 21 insertions(+), 10 deletions(-)

diff --git a/trezorctl b/trezorctl
index 1f7a10d78a..52d82050c2 100755
--- a/trezorctl
+++ b/trezorctl
@@ -28,7 +28,7 @@ import logging
 import os
 import sys
 
-from trezorlib.client import TrezorClient, TrezorClientVerbose, CallException
+from trezorlib.client import TrezorClient, CallException
 from trezorlib.transport import get_transport, enumerate_devices
 from trezorlib import coins
 from trezorlib import log
@@ -67,11 +67,7 @@ CHOICE_OUTPUT_SCRIPT_TYPE = ChoiceType({
 
 
 def enable_logging():
-    handler = logging.StreamHandler()
-    handler.setFormatter(log.PrettyProtobufFormatter())
-    logger = logging.getLogger('trezorlib')
-    logger.setLevel(logging.DEBUG)
-    logger.addHandler(handler)
+    log.enable_debug_output()
     log.OMITTED_MESSAGES.add(proto.Features)
 
 
diff --git a/trezorlib/log.py b/trezorlib/log.py
index f23e39ed3a..0e07232fa3 100644
--- a/trezorlib/log.py
+++ b/trezorlib/log.py
@@ -1,5 +1,5 @@
 import logging
-from typing import Set, Type
+from typing import Set, Type, Optional
 
 from . import protobuf
 
@@ -8,12 +8,27 @@ OMITTED_MESSAGES = set()  # type: Set[Type[protobuf.MessageType]]
 
 class PrettyProtobufFormatter(logging.Formatter):
 
-    def format(self, record):
+    def format(self, record: logging.LogRecord) -> str:
         time = self.formatTime(record)
-        message = '[{time}] {level}: {msg}'.format(time=time, level=record.levelname, msg=super().format(record))
+        message = '[{time}] {source} {level}: {msg}'.format(
+            time=time, level=record.levelname.upper(),
+            source=record.name,
+            msg=super().format(record))
         if hasattr(record, 'protobuf'):
             if type(record.protobuf) in OMITTED_MESSAGES:
                 message += " ({} bytes)".format(record.protobuf.ByteSize())
             else:
-                message += "\n" + protobuf.format_message(record.protobuf)
+                message += "\n" + protobuf.format_message(record.protobuf, collapse_cointypes=True)
         return message
+
+
+def enable_debug_output(handler: Optional[logging.Handler] = None):
+    if handler is None:
+        handler = logging.StreamHandler()
+
+    formatter = PrettyProtobufFormatter()
+    handler.setFormatter(formatter)
+
+    logger = logging.getLogger('trezorlib')
+    logger.setLevel(logging.DEBUG)
+    logger.addHandler(handler)