Add -P,--cherry-pick option to "b4 am"

This lets someone select a subset of patches in a series, e.g.:

b4 am -P 1-3,5,7- [msgid]

Suggested-by: Heiko Stübner <heiko@sntech.de>
Signed-off-by: Konstantin Ryabitsev <konstantin@linuxfoundation.org>

3 files changed, 31 insertions, 13 deletions

diff --git a/b4/__init__.py b/b4/__init__.py
index c9011ac..e7eab62 100644
--- a/b4/__init__.py
+++ b/b4/__init__.py
@@ -465,8 +465,8 @@ class LoreSeries:
 
         return slug
 
-    def save_am_mbox(self, mbx, noaddtrailers, covertrailers,
-                     trailer_order=None, addmysob=False, addlink=False, linkmask=None):
+    def save_am_mbox(self, mbx, noaddtrailers, covertrailers, trailer_order=None, addmysob=False,
+                     addlink=False, linkmask=None, cherrypick=None):
 
         usercfg = get_user_config()
         config = get_main_config()
@@ -496,6 +496,10 @@ class LoreSeries:
         at = 1
         atterrors = list()
         for lmsg in self.patches[1:]:
+            if cherrypick is not None and at not in cherrypick:
+                at += 1
+                logger.debug('  skipped: [%s/%s] (not in cherrypick)', at, self.expected)
+                continue
             if lmsg is not None:
                 if self.has_cover and covertrailers and self.patches[0].followup_trailers:
                     lmsg.followup_trailers.update(self.patches[0].followup_trailers)
@@ -1870,7 +1874,7 @@ def parse_int_range(intrange, upper=None):
         if n.isdigit():
             yield int(n)
         elif n.find('<') == 0 and len(n) > 1 and n[1:].isdigit():
-            yield from range(0, int(n[1:]))
+            yield from range(1, int(n[1:]))
         elif n.find('-') > 0:
             nr = n.split('-')
             if nr[0].isdigit() and nr[1].isdigit():
diff --git a/b4/command.py b/b4/command.py
index f83c935..097de5f 100644
--- a/b4/command.py
+++ b/b4/command.py
@@ -95,6 +95,8 @@ def cmd():
                        help='Add a lore.kernel.org/r/ link to every patch')
     sp_am.add_argument('-Q', '--quilt-ready', dest='quiltready', action='store_true', default=False,
                        help='Save mbox patches in a quilt-ready folder')
+    sp_am.add_argument('-P', '--cherry-pick', dest='cherrypick', default=None,
+                       help='Cherry-pick a subset of patches (e.g. -P 1-2,4,6-)')
     sp_am.set_defaults(func=cmd_am)
 
     # b4 attest
diff --git a/b4/mbox.py b/b4/mbox.py
index f1b4700..1ffb563 100644
--- a/b4/mbox.py
+++ b/b4/mbox.py
@@ -66,15 +66,22 @@ def mbox_to_am(mboxfile, cmdargs):
         os.unlink(am_filename)
 
     logger.info('---')
+    if cmdargs.cherrypick:
+        cherrypick = list(b4.parse_int_range(cmdargs.cherrypick, upper=len(lser.patches)-1))
+    else:
+        cherrypick = None
     logger.critical('Writing %s', am_filename)
     mbx = mailbox.mbox(am_filename)
     am_mbx = lser.save_am_mbox(mbx, cmdargs.noaddtrailers, covertrailers,
                                trailer_order=config['trailer-order'],
                                addmysob=cmdargs.addmysob, addlink=cmdargs.addlink,
-                               linkmask=config['linkmask'])
+                               linkmask=config['linkmask'], cherrypick=cherrypick)
     logger.info('---')
 
-    logger.critical('Total patches: %s', len(am_mbx))
+    if cherrypick is None:
+        logger.critical('Total patches: %s', len(am_mbx))
+    else:
+        logger.info('Total patches: %s (cherrypicked: %s)', len(am_mbx), cmdargs.cherrypick)
     if lser.has_cover and lser.patches[0].followup_trailers and not covertrailers:
         # Warn that some trailers were sent to the cover letter
         logger.critical('---')
@@ -136,12 +143,12 @@ def mbox_to_am(mboxfile, cmdargs):
         logger.critical('       git am %s', am_filename)
 
     am_mbx.close()
-    thanks_record_am(lser)
+    thanks_record_am(lser, cherrypick=cherrypick)
 
     return am_filename
 
 
-def thanks_record_am(lser):
+def thanks_record_am(lser, cherrypick=None):
     if not lser.complete:
         logger.debug('Incomplete series, not tracking for thanks')
         return
@@ -150,13 +157,18 @@ def thanks_record_am(lser):
     datadir = b4.get_data_dir()
     slug = lser.get_slug(extended=True)
     filename = '%s.am' % slug
-    # Check if we're tracking it already
-    for entry in os.listdir(datadir):
-        if entry == filename:
-            return
 
     patches = list()
+    at = 0
     for pmsg in lser.patches[1:]:
+        at += 1
+        if pmsg is None:
+            continue
+
+        if cherrypick is not None and at not in cherrypick:
+            logger.debug('Skipped non-cherrypicked: %s', at)
+            continue
+
         pmsg.load_hashes()
         if pmsg.attestation is None:
             logger.debug('Unable to get hashes for all patches, not tracking for thanks')
@@ -167,8 +179,8 @@ def thanks_record_am(lser):
     if lmsg is None:
         lmsg = lser.patches[1]
 
-    allto = email.utils.getaddresses([b4.LoreMessage.clean_header(x) for x in lmsg.msg.get_all('to', [])])
-    allcc = email.utils.getaddresses([b4.LoreMessage.clean_header(x) for x in lmsg.msg.get_all('cc', [])])
+    allto = email.utils.getaddresses([str(x) for x in lmsg.msg.get_all('to', [])])
+    allcc = email.utils.getaddresses([str(x) for x in lmsg.msg.get_all('cc', [])])
 
     out = {
         'msgid': lmsg.msgid,