Add "b4 shazam" that is like b4 am + git am

By popular demand, provide a way to apply series straight to a git
repository. By default, we're still going the safest possible route:

- create a sparse worktree consisting just of the files being modified
- run "git am" against the temporary worktree
- if "git am" went well, fetch from the temporary worktree into our
  current tree and leave everything in FETCH_HEAD
- unless we're running "b4 shazam -A" in which case we just apply to the
  current HEAD (exact equivalent of b4 am -o- | git am)

Further changes to come based on feedback.

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

3 files changed, 210 insertions, 124 deletions

diff --git a/b4/__init__.py b/b4/__init__.py
index bc7b8dd..47272bb 100644
--- a/b4/__init__.py
+++ b/b4/__init__.py
@@ -44,7 +44,7 @@ try:
 except ModuleNotFoundError:
     can_patatt = False
 
-__VERSION__ = '0.8.0'
+__VERSION__ = '0.9-dev'
 
 
 def _dkim_log_filter(record):
@@ -421,8 +421,7 @@ class LoreSeries:
         self.has_cover = False
         self.partial_reroll = False
         self.subject = '(untitled)'
-        # Used for base matching
-        self._indexes = None
+        self.indexes = None
 
     def __repr__(self):
         out = list()
@@ -601,28 +600,31 @@ class LoreSeries:
 
         return msgs
 
-    def check_applies_clean(self, gitdir: str, at: Optional[str] = None) -> Tuple[int, list]:
-        if self._indexes is None:
-            self._indexes = list()
-            seenfiles = set()
-            for lmsg in self.patches[1:]:
-                if lmsg is None or lmsg.blob_indexes is None:
+    def populate_indexes(self):
+        self.indexes = list()
+        seenfiles = set()
+        for lmsg in self.patches[1:]:
+            if lmsg is None or lmsg.blob_indexes is None:
+                continue
+            for fn, bh in lmsg.blob_indexes:
+                if fn in seenfiles:
+                    # if we have seen this file once already, then it's a repeat patch
+                    # and it's no longer going to match current hash
                     continue
-                for fn, bh in lmsg.blob_indexes:
-                    if fn in seenfiles:
-                        # if we have seen this file once already, then it's a repeat patch
-                        # and it's no longer going to match current hash
-                        continue
-                    seenfiles.add(fn)
-                    if set(bh) == {'0'}:
-                        # New file, will for sure apply clean
-                        continue
-                    self._indexes.append((fn, bh))
+                seenfiles.add(fn)
+                if set(bh) == {'0'}:
+                    # New file, will for sure apply clean
+                    continue
+                self.indexes.append((fn, bh))
+
+    def check_applies_clean(self, gitdir: str, at: Optional[str] = None) -> Tuple[int, list]:
+        if self.indexes is None:
+            self.populate_indexes()
 
         mismatches = list()
         if at is None:
             at = 'HEAD'
-        for fn, bh in self._indexes:
+        for fn, bh in self.indexes:
             ecode, out = git_run_command(gitdir, ['ls-tree', at, fn])
             if ecode == 0 and len(out):
                 chunks = out.split()
@@ -636,7 +638,7 @@ class LoreSeries:
                 logger.debug('Could not look up %s:%s', at, fn)
             mismatches.append((fn, bh))
 
-        return len(self._indexes), mismatches
+        return len(self.indexes), mismatches
 
     def find_base(self, gitdir: str, branches: Optional[str] = None, maxdays: int = 30) -> Tuple[str, len, len]:
         # Find the date of the first patch we have
@@ -695,13 +697,13 @@ class LoreSeries:
                     break
         else:
             best = commit
-        if fewest == len(self._indexes):
+        if fewest == len(self.indexes):
             # None of the blobs matched
             raise IndexError
 
         lines = git_get_command_lines(gitdir, ['describe', '--all', best])
         if len(lines):
-            return lines[0], len(self._indexes), fewest
+            return lines[0], len(self.indexes), fewest
 
         raise IndexError
 
diff --git a/b4/command.py b/b4/command.py
index 5bb3384..199d0c2 100644
--- a/b4/command.py
+++ b/b4/command.py
@@ -17,7 +17,7 @@ def cmd_retrieval_common_opts(sp):
     sp.add_argument('msgid', nargs='?',
                     help='Message ID to process, or pipe a raw message')
     sp.add_argument('-p', '--use-project', dest='useproject', default=None,
-                    help='Use a specific project instead of guessing (linux-mm, linux-hardening, etc)')
+                    help='Use a specific project instead of default (linux-mm, linux-hardening, etc)')
     sp.add_argument('-m', '--use-local-mbox', dest='localmbox', default=None,
                     help='Instead of grabbing a thread from lore, process this mbox file (or - for stdin)')
     sp.add_argument('-C', '--no-cache', dest='nocache', action='store_true', default=False,
@@ -35,6 +35,26 @@ def cmd_mbox_common_opts(sp):
     sp.add_argument('-M', '--save-as-maildir', dest='maildir', action='store_true', default=False,
                     help='Save as maildir (avoids mbox format ambiguities)')
 
+def cmd_am_common_opts(sp):
+    sp.add_argument('-v', '--use-version', dest='wantver', type=int, default=None,
+                    help='Get a specific version of the patch/series')
+    sp.add_argument('-t', '--apply-cover-trailers', dest='covertrailers', action='store_true', default=False,
+                    help='Apply trailers sent to the cover letter to all patches')
+    sp.add_argument('-S', '--sloppy-trailers', dest='sloppytrailers', action='store_true', default=False,
+                    help='Apply trailers without email address match checking')
+    sp.add_argument('-T', '--no-add-trailers', dest='noaddtrailers', action='store_true', default=False,
+                    help='Do not add or sort any trailers')
+    sp.add_argument('-s', '--add-my-sob', dest='addmysob', action='store_true', default=False,
+                    help='Add your own signed-off-by to every patch')
+    sp.add_argument('-l', '--add-link', dest='addlink', action='store_true', default=False,
+                    help='Add a Link: with message-id lookup URL to every patch')
+    sp.add_argument('-P', '--cherry-pick', dest='cherrypick', default=None,
+                    help='Cherry-pick a subset of patches (e.g. "-P 1-2,4,6-", '
+                         '"-P _" to use just the msgid specified, or '
+                         '"-P *globbing*" to match on commit subject)')
+    sp.add_argument('--cc-trailers', dest='copyccs', action='store_true', default=False,
+                    help='Copy all Cc\'d addresses into Cc: trailers')
+
 
 def cmd_mbox(cmdargs):
     import b4.mbox
@@ -51,6 +71,11 @@ def cmd_am(cmdargs):
     b4.mbox.main(cmdargs)
 
 
+def cmd_shazam(cmdargs):
+    import b4.mbox
+    b4.mbox.main(cmdargs)
+
+
 def cmd_attest(cmdargs):
     import b4.attest
     if len(cmdargs.patchfile):
@@ -100,41 +125,34 @@ def cmd():
     # b4 am
     sp_am = subparsers.add_parser('am', help='Create an mbox file that is ready to git-am')
     cmd_mbox_common_opts(sp_am)
-    sp_am.add_argument('-v', '--use-version', dest='wantver', type=int, default=None,
-                       help='Get a specific version of the patch/series')
-    sp_am.add_argument('-t', '--apply-cover-trailers', dest='covertrailers', action='store_true', default=False,
-                       help='Apply trailers sent to the cover letter to all patches')
-    sp_am.add_argument('-S', '--sloppy-trailers', dest='sloppytrailers', action='store_true', default=False,
-                       help='Apply trailers without email address match checking')
-    sp_am.add_argument('-T', '--no-add-trailers', dest='noaddtrailers', action='store_true', default=False,
-                       help='Do not add or sort any trailers')
-    sp_am.add_argument('-s', '--add-my-sob', dest='addmysob', action='store_true', default=False,
-                       help='Add your own signed-off-by to every patch')
-    sp_am.add_argument('-l', '--add-link', dest='addlink', action='store_true', default=False,
-                       help='Add a lore.kernel.org/r/ link to every patch')
+    cmd_am_common_opts(sp_am)
     sp_am.add_argument('-Q', '--quilt-ready', dest='quiltready', action='store_true', default=False,
                        help='Save 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-", '
-                            '"-P _" to use just the msgid specified, or '
-                            '"-P *globbing*" to match on commit subject)')
     sp_am.add_argument('-g', '--guess-base', dest='guessbase', action='store_true', default=False,
                        help='Try to guess the base of the series (if not specified)')
     sp_am.add_argument('-b', '--guess-branch', dest='guessbranch', default=None,
                        help='When guessing base, restrict to this branch (use with -g)')
-    sp_am.add_argument('--guess-lookback', dest='guessdays', type=int, default=14,
-                       help='When guessing base, go back this many days from the date of the patch')
+    sp_am.add_argument('--guess-lookback', dest='guessdays', type=int, default=21,
+                       help='When guessing base, go back this many days from the patch date (default: 2 weeks)')
     sp_am.add_argument('-3', '--prep-3way', dest='threeway', action='store_true', default=False,
                        help='Prepare for a 3-way merge '
                             '(tries to ensure that all index blobs exist by making a fake commit range)')
-    sp_am.add_argument('--cc-trailers', dest='copyccs', action='store_true', default=False,
-                       help='Copy all Cc\'d addresses into Cc: trailers')
     sp_am.add_argument('--no-cover', dest='nocover', action='store_true', default=False,
                        help='Do not save the cover letter (on by default when using -o -)')
     sp_am.add_argument('--no-partial-reroll', dest='nopartialreroll', action='store_true', default=False,
                        help='Do not reroll partial series when detected')
     sp_am.set_defaults(func=cmd_am)
 
+    # b4 shazam
+    sp_sh = subparsers.add_parser('shazam', help='Like b4 am, but applies the series to your tree')
+    cmd_retrieval_common_opts(sp_sh)
+    cmd_am_common_opts(sp_sh)
+    sp_sh.add_argument('--guess-lookback', dest='guessdays', type=int, default=21,
+                       help = 'When guessing base, go back this many days from the patch date (default: 3 weeks)')
+    sp_sh.add_argument('-A', '--apply-here', dest='applyhere', action='store_true', default=False,
+                       help = 'Apply to the current tree')
+    sp_sh.set_defaults(func=cmd_shazam)
+
     # b4 attest
     sp_att = subparsers.add_parser('attest', help='Create cryptographic attestation for a set of patches')
     sp_att.add_argument('-f', '--from', dest='sender', default=None,
@@ -200,7 +218,7 @@ def cmd():
     sp_diff.add_argument('-g', '--gitdir', default=None,
                          help='Operate on this git tree instead of current dir')
     sp_diff.add_argument('-p', '--use-project', dest='useproject', default=None,
-                         help='Use a specific project instead of guessing (linux-mm, linux-hardening, etc)')
+                         help='Use a specific project instead of default (linux-mm, linux-hardening, etc)')
     sp_diff.add_argument('-C', '--no-cache', dest='nocache', action='store_true', default=False,
                          help='Do not use local cache')
     sp_diff.add_argument('-v', '--compare-versions', dest='wantvers', type=int, default=None, nargs='+',
diff --git a/b4/mbox.py b/b4/mbox.py
index a7e9f78..3e3ae05 100644
--- a/b4/mbox.py
+++ b/b4/mbox.py
@@ -18,6 +18,7 @@ import fnmatch
 import shutil
 import pathlib
 import tempfile
+import io
 
 import urllib.parse
 import xml.etree.ElementTree
@@ -35,7 +36,6 @@ def make_am(msgs, cmdargs, msgid):
     if outdir == '-':
         cmdargs.nocover = True
     wantver = cmdargs.wantver
-    wantname = cmdargs.wantname
     covertrailers = cmdargs.covertrailers
     count = len(msgs)
     logger.info('Analyzing %s messages in the thread', count)
@@ -97,41 +97,6 @@ def make_am(msgs, cmdargs, msgid):
     except KeyError:
         sys.exit(1)
 
-    if cmdargs.maildir or config.get('save-maildirs', 'no') == 'yes':
-        save_maildir = True
-        dftext = 'maildir'
-    else:
-        save_maildir = False
-        dftext = 'mbx'
-
-    if wantname:
-        slug = wantname
-        if wantname.find('.') > -1:
-            slug = '.'.join(wantname.split('.')[:-1])
-        gitbranch = slug
-    else:
-        slug = lser.get_slug(extended=True)
-        gitbranch = lser.get_slug(extended=False)
-
-    if outdir != '-':
-        am_filename = os.path.join(outdir, f'{slug}.{dftext}')
-        am_cover = os.path.join(outdir, f'{slug}.cover')
-
-        if os.path.exists(am_filename):
-            if os.path.isdir(am_filename):
-                shutil.rmtree(am_filename)
-            else:
-                os.unlink(am_filename)
-        if save_maildir:
-            b4.save_maildir(am_msgs, am_filename)
-        else:
-            with open(am_filename, 'w') as fh:
-                b4.save_git_am_mbox(am_msgs, fh)
-    else:
-        am_filename = None
-        am_cover = None
-        b4.save_git_am_mbox(am_msgs, sys.stdout)
-
     logger.info('---')
 
     if cherrypick is None:
@@ -173,6 +138,17 @@ def make_am(msgs, cmdargs, msgid):
             logger.critical('     Msg From: %s <%s>', fname, femail)
         logger.critical('NOTE: Rerun with -S to apply them anyway')
 
+    top_msgid = None
+    first_body = None
+    for lmsg in lser.patches:
+        if lmsg is not None:
+            first_body = lmsg.body
+            top_msgid = lmsg.msgid
+            break
+    if top_msgid is None:
+        logger.critical('Could not find any patches in the series.')
+        return
+
     topdir = b4.git_get_toplevel()
 
     if cmdargs.threeway:
@@ -194,27 +170,54 @@ def make_am(msgs, cmdargs, msgid):
     if not lser.complete and not cmdargs.cherrypick:
         logger.critical('WARNING: Thread incomplete!')
 
-    if lser.has_cover and not cmdargs.nocover:
-        lser.save_cover(am_cover)
+    gitbranch = lser.get_slug(extended=False)
+    am_filename = None
 
-    top_msgid = None
-    first_body = None
-    for lmsg in lser.patches:
-        if lmsg is not None:
-            first_body = lmsg.body
-            top_msgid = lmsg.msgid
-            break
-    if top_msgid is None:
-        logger.critical('Could not find any patches in the series.')
-        return
+    if cmdargs.subcmd == 'am':
+        wantname = cmdargs.wantname
+        if cmdargs.maildir or config.get('save-maildirs', 'no') == 'yes':
+            save_maildir = True
+            dftext = 'maildir'
+        else:
+            save_maildir = False
+            dftext = 'mbx'
+
+        if wantname:
+            slug = wantname
+            if wantname.find('.') > -1:
+                slug = '.'.join(wantname.split('.')[:-1])
+            gitbranch = slug
+        else:
+            slug = lser.get_slug(extended=True)
+
+        if outdir != '-':
+            am_filename = os.path.join(outdir, f'{slug}.{dftext}')
+            am_cover = os.path.join(outdir, f'{slug}.cover')
 
-    linkurl = config['linkmask'] % top_msgid
-    if cmdargs.quiltready:
-        q_dirname = os.path.join(outdir, f'{slug}.patches')
-        save_as_quilt(am_msgs, q_dirname)
-        logger.critical('Quilt: %s', q_dirname)
+            if os.path.exists(am_filename):
+                if os.path.isdir(am_filename):
+                    shutil.rmtree(am_filename)
+                else:
+                    os.unlink(am_filename)
+            if save_maildir:
+                b4.save_maildir(am_msgs, am_filename)
+            else:
+                with open(am_filename, 'w') as fh:
+                    b4.save_git_am_mbox(am_msgs, fh)
+        else:
+            am_cover = None
+            b4.save_git_am_mbox(am_msgs, sys.stdout)
 
-    logger.critical(' Link: %s', linkurl)
+        if lser.has_cover and not cmdargs.nocover:
+            lser.save_cover(am_cover)
+
+        linkurl = config['linkmask'] % top_msgid
+        if cmdargs.quiltready:
+            q_dirname = os.path.join(outdir, f'{slug}.patches')
+            save_as_quilt(am_msgs, q_dirname)
+            logger.critical('Quilt: %s', q_dirname)
+
+        logger.critical(' Link: %s', linkurl)
 
     base_commit = None
     matches = re.search(r'base-commit: .*?([0-9a-f]+)', first_body, re.MULTILINE)
@@ -226,30 +229,82 @@ def make_am(msgs, cmdargs, msgid):
         if matches:
             base_commit = matches.groups()[0]
 
-    if base_commit:
-        logger.critical(' Base: %s', base_commit)
-    else:
-        if topdir is not None:
-            if cmdargs.guessbase:
-                logger.critical('       attempting to guess base-commit...')
-                try:
-                    base_commit, nblobs, mismatches = lser.find_base(topdir, branches=cmdargs.guessbranch,
-                                                                     maxdays=cmdargs.guessdays)
-                    if mismatches == 0:
-                        logger.critical(' Base: %s (exact match)', base_commit)
-                    elif nblobs == mismatches:
-                        logger.critical(' Base: failed to guess base')
-                    else:
-                        logger.critical(' Base: %s (best guess, %s/%s blobs matched)', base_commit,
-                                        nblobs - mismatches, nblobs)
-                except IndexError:
-                    logger.critical(' Base: failed to guess base')
+    if not base_commit and topdir and cmdargs.guessbase:
+        logger.critical(' Base: attempting to guess base-commit...')
+        try:
+            base_commit, nblobs, mismatches = lser.find_base(topdir, branches=cmdargs.guessbranch,
+                                                             maxdays=cmdargs.guessdays)
+            if mismatches == 0:
+                logger.critical(' Base: %s (exact match)', base_commit)
+            elif nblobs == mismatches:
+                logger.critical(' Base: failed to guess base')
             else:
-                checked, mismatches = lser.check_applies_clean(topdir, at=cmdargs.guessbranch)
-                if checked and len(mismatches) == 0 and checked != mismatches:
-                    logger.critical(' Base: applies clean to current tree')
-                else:
-                    logger.critical(' Base: not specified')
+                logger.critical(' Base: %s (best guess, %s/%s blobs matched)', base_commit,
+                                nblobs - mismatches, nblobs)
+        except IndexError:
+            logger.critical(' Base: failed to guess base')
+
+    if cmdargs.subcmd == 'shazam':
+        if not topdir:
+            logger.critical('Could not figure out where your git dir is, cannot shazam.')
+            sys.exit(1)
+        ifh = io.StringIO()
+        b4.save_git_am_mbox(am_msgs, ifh)
+        ambytes = ifh.getvalue().encode()
+        if cmdargs.applyhere:
+            # Blindly attempt to apply to the current tree
+            ecode, out = b4.git_run_command(topdir, ['am'], stdin=ambytes, logstderr=True)
+            logger.info(out.strip())
+            sys.exit(ecode)
+
+        if not base_commit:
+            # Try our best with HEAD, I guess
+            base_commit = 'HEAD'
+
+        with b4.git_temp_worktree(topdir, base_commit) as gwt:
+            if lser.indexes is None:
+                lser.populate_indexes()
+            # TODO: Handle patches containing nothing but new file additions
+            wantfiles = [i[0] for i in lser.indexes]
+            logger.info('Magic: Preparing a sparse worktree with %s files', len(wantfiles))
+            # TODO: Handle these potential errors
+            ecode, out = b4.git_run_command(gwt, ['sparse-checkout', 'init'] + wantfiles, logstderr=True)
+            if ecode > 0:
+                logger.critical('Error running sparse-checkout init')
+                logger.critical(out)
+                sys.exit(ecode)
+            ecode, out = b4.git_run_command(gwt, ['checkout'], logstderr=True)
+            if ecode > 0:
+                logger.critical('Error running checkout into sparse workdir')
+                logger.critical(out)
+                sys.exit(ecode)
+            ecode, out = b4.git_run_command(gwt, ['am'], stdin=ambytes, logstderr=True)
+            if ecode > 0:
+                logger.critical('Unable to cleanly apply series, see failure log below')
+                logger.critical('---')
+                logger.critical(out.strip())
+                logger.critical('---')
+                logger.critical('Not fetching into FETCH_HEAD')
+                sys.exit(ecode)
+            logger.info('---')
+            logger.info(out.strip())
+            logger.info('---')
+            logger.info('Fetching into FETCH_HEAD')
+            gitargs = ['fetch', gwt]
+            ecode, out = b4.git_run_command(topdir, gitargs, logstderr=True)
+            if ecode > 0:
+                logger.critical('Unable to fetch from the worktree')
+                logger.critical(out.strip())
+                sys.exit(ecode)
+        logger.info('You can now merge or checkout FETCH_HEAD')
+        thanks_record_am(lser, cherrypick=cherrypick)
+        return
+
+    if not base_commit:
+        checked, mismatches = lser.check_applies_clean(topdir, at=cmdargs.guessbranch)
+        if checked and len(mismatches) == 0 and checked != mismatches:
+            logger.critical(' Base: applies clean to current tree')
+            base_commit = 'HEAD'
         else:
             logger.critical(' Base: not specified')
 
@@ -258,7 +313,6 @@ def make_am(msgs, cmdargs, msgid):
     if cmdargs.outdir != '-':
         logger.critical('       git am %s', am_filename)
 
-    thanks_record_am(lser, cherrypick=cherrypick)
 
 
 def thanks_record_am(lser, cherrypick=None):
@@ -578,6 +632,18 @@ def get_msgs(cmdargs) -> Tuple[Optional[str], Optional[list]]:
 
 
 def main(cmdargs):
+    if cmdargs.subcmd == 'shazam':
+        # We force some settings
+        cmdargs.checknewer = True
+        cmdargs.threeway = False
+        cmdargs.nopartialreroll = False
+        cmdargs.outdir = '-'
+        cmdargs.guessbranch = None
+        if cmdargs.applyhere:
+            cmdargs.guessbase = False
+        else:
+            cmdargs.guessbase = True
+
     if cmdargs.checknewer:
         # Force nocache mode
         cmdargs.nocache = True
@@ -589,7 +655,7 @@ def main(cmdargs):
     if len(msgs) and cmdargs.checknewer:
         msgs = get_extra_series(msgs, direction=1, useproject=cmdargs.useproject)
 
-    if cmdargs.subcmd == 'am':
+    if cmdargs.subcmd in ('am', 'shazam'):
         make_am(msgs, cmdargs, msgid)
         return