Save exploded pull requests as maildirs as well

This moves maildir saving code into __init__.py so that we can benefit
from it via other subcommands, such as pr.

Signed-off-by: Konstantin Ryabitsev <konstantin@linuxfoundation.org>

3 files changed, 36 insertions, 22 deletions

diff --git a/b4/__init__.py b/b4/__init__.py
index 2ab2d4b..cad22b7 100644
--- a/b4/__init__.py
+++ b/b4/__init__.py
@@ -605,8 +605,7 @@ class LoreSeries:
                 if noaddtrailers:
                     add_trailers = False
                 msg = lmsg.get_am_message(add_trailers=add_trailers, trailer_order=trailer_order, copyccs=copyccs)
-                slug = '%04d_%s' % (lmsg.counter, re.sub(r'\W+', '_', lmsg.subject).strip('_').lower())
-                msgs.append((slug, msg))
+                msgs.append(msg)
             else:
                 logger.error('  ERROR: missing [%s/%s]!', at, self.expected)
             at += 1
@@ -2330,12 +2329,28 @@ def save_git_am_mbox(msgs: list, dest):
     # So, save in the format that git-am expects
     # "dest" should be a file handler in writable+binary mode
     for msg in msgs:
-        bmsg = msg.as_bytes(unixfrom=True, policy=emlpolicy)
+        bmsg = msg.as_string(unixfrom=True, policy=emlpolicy)
         # public-inbox unixfrom says "mboxrd", so replace it with something else
         # so there is no confusion as it's NOT mboxrd
-        bmsg = re.sub(b'^From mboxrd@z ', b'From git@z ', bmsg)
-        bmsg = bmsg.rstrip(b'\r\n') + b'\n\n'
-        dest.write(bmsg)
+        bmsg = re.sub('^From mboxrd@z ', 'From git@z ', bmsg)
+        bmsg = bmsg.rstrip('\r\n') + '\n\n'
+        dest.write(bmsg.encode())
+
+
+def save_maildir(msgs: list, dest):
+    d_new = os.path.join(dest, 'new')
+    pathlib.Path(d_new).mkdir(parents=True)
+    d_cur = os.path.join(dest, 'cur')
+    pathlib.Path(d_cur).mkdir(parents=True)
+    d_tmp = os.path.join(dest, 'tmp')
+    pathlib.Path(d_tmp).mkdir(parents=True)
+    for msg in msgs:
+        # make a slug out of it
+        lsubj = LoreSubject(msg.get('subject', ''))
+        slug = '%04d_%s' % (lsubj.counter, re.sub(r'\W+', '_', lsubj.subject).strip('_').lower())
+        with open(os.path.join(d_tmp, f'{slug}.eml'), 'wb') as mfh:
+            mfh.write(msg.as_string(policy=emlpolicy).encode())
+        os.rename(os.path.join(d_tmp, f'{slug}.eml'), os.path.join(d_new, f'{slug}.eml'))
 
 
 def get_lore_projects_from_msg(msg) -> list:
diff --git a/b4/mbox.py b/b4/mbox.py
index 33d7772..c0ada06 100644
--- a/b4/mbox.py
+++ b/b4/mbox.py
@@ -123,19 +123,10 @@ def make_am(msgs, cmdargs, msgid):
             else:
                 os.unlink(am_filename)
         if save_maildir:
-            d_new = os.path.join(am_filename, 'new')
-            pathlib.Path(d_new).mkdir(parents=True)
-            d_cur = os.path.join(am_filename, 'cur')
-            pathlib.Path(d_cur).mkdir(parents=True)
-            d_tmp = os.path.join(am_filename, 'tmp')
-            pathlib.Path(d_tmp).mkdir(parents=True)
-            for m_slug, msg in am_msgs:
-                with open(os.path.join(d_tmp, f'{m_slug}.eml'), 'wb') as mfh:
-                    mfh.write(msg.as_bytes(policy=b4.emlpolicy))
-                os.rename(os.path.join(d_tmp, f'{m_slug}.eml'), os.path.join(d_new, f'{m_slug}.eml'))
+            b4.save_maildir(am_msgs, am_filename)
         else:
             with open(am_filename, 'wb') as fh:
-                b4.save_git_am_mbox([x[1] for x in am_msgs], fh)
+                b4.save_git_am_mbox(am_msgs, fh)
     else:
         am_filename = None
         am_cover = None
diff --git a/b4/pr.py b/b4/pr.py
index 86d039c..98de07b 100644
--- a/b4/pr.py
+++ b/b4/pr.py
@@ -456,9 +456,16 @@ def main(cmdargs):
         lmsg.pr_tip_commit = lmsg.pr_remote_tip_commit
 
     if cmdargs.explode:
+        config = b4.get_main_config()
+        if config.get('save-maildirs', 'no') == 'yes':
+            save_maildir = True
+            dftext = 'maildir'
+        else:
+            save_maildir = False
+            dftext = 'mbx'
         savefile = cmdargs.outmbox
         if savefile is None:
-            savefile = '%s.mbx' % lmsg.msgid
+            savefile = f'{lmsg.msgid}.{dftext}'
         if os.path.exists(savefile):
             logger.info('File exists: %s', savefile)
             sys.exit(1)
@@ -472,10 +479,11 @@ def main(cmdargs):
                 sys.exit(1)
 
             if msgs:
-                smbx = mailbox.mbox(savefile)
-                for msg in msgs:
-                    smbx.add(msg.as_string(policy=b4.emlpolicy).encode())
-                smbx.close()
+                if save_maildir:
+                    b4.save_maildir(msgs, savefile)
+                else:
+                    with open(savefile, 'wb') as fh:
+                        b4.save_git_am_mbox(msgs, fh)
                 logger.info('---')
                 logger.info('Saved %s', savefile)
                 sys.exit(0)