python: refactor print_messages into format_messages and print_messages

1 files changed, 26 insertions, 11 deletions

diff --git a/bindings/python/notmuch/message.py b/bindings/python/notmuch/message.py
index 245e8148..fa6c58c3 100644
--- a/bindings/python/notmuch/message.py
+++ b/bindings/python/notmuch/message.py
@@ -188,14 +188,17 @@ class Messages(object):
         if self._msgs is not None:
             self._destroy(self._msgs)
 
-    def print_messages(self, format, indent=0, entire_thread=False):
-        """Outputs messages as needed for 'notmuch show' to sys.stdout
+    def format_messages(self, format, indent=0, entire_thread=False):
+        """Formats messages as needed for 'notmuch show'.
 
         :param format: A string of either 'text' or 'json'.
         :param indent: A number indicating the reply depth of these messages.
         :param entire_thread: A bool, indicating whether we want to output
                        whole threads or only the matching messages.
+        :return: a list of lines
         """
+        result = list()
+
         if format.lower() == "text":
             set_start = ""
             set_end = ""
@@ -209,36 +212,48 @@ class Messages(object):
 
         first_set = True
 
-        sys.stdout.write(set_start)
+        result.append(set_start)
 
         # iterate through all toplevel messages in this thread
         for msg in self:
             # if not msg:
             #     break
             if not first_set:
-                sys.stdout.write(set_sep)
+                result.append(set_sep)
             first_set = False
 
-            sys.stdout.write(set_start)
+            result.append(set_start)
             match = msg.is_match()
             next_indent = indent
 
             if (match or entire_thread):
                 if format.lower() == "text":
-                    sys.stdout.write(msg.format_message_as_text(indent))
+                    result.append(msg.format_message_as_text(indent))
                 else:
-                    sys.stdout.write(msg.format_message_as_json(indent))
+                    result.append(msg.format_message_as_json(indent))
                 next_indent = indent + 1
 
             # get replies and print them also out (if there are any)
             replies = msg.get_replies()
             if not replies is None:
-                sys.stdout.write(set_sep)
-                replies.print_messages(format, next_indent, entire_thread)
+                result.append(set_sep)
+                result.extend(replies.format_messages(format, next_indent, entire_thread))
+
+            result.append(set_end)
+        result.append(set_end)
 
-            sys.stdout.write(set_end)
-        sys.stdout.write(set_end)
+        return result
 
+    def print_messages(self, format, indent=0, entire_thread=False, handle=sys.stdout):
+        """Outputs messages as needed for 'notmuch show' to a file like object.
+
+        :param format: A string of either 'text' or 'json'.
+        :param handle: A file like object to print to (default is sys.stdout).
+        :param indent: A number indicating the reply depth of these messages.
+        :param entire_thread: A bool, indicating whether we want to output
+                       whole threads or only the matching messages.
+        """
+        handle.write(''.join(self.format_messages(format, indent, entire_thread)))
 
 class Message(Python3StringMixIn):
     """Represents a single Email message