hgext/hgshelve.py

   1 # shelve.py
   2 #
   3 # Copyright 2007 Bryan O'Sullivan <bos@serpentine.com>
   4 # Copyright 2007 TK Soh <teekaysoh@gmailcom>
   5 #
   6 # This software may be used and distributed according to the terms of
   7 # the GNU General Public License, incorporated herein by reference.
   8
   9 '''interactive change selection to set aside that may be restored later'''
  10
  11 from mercurial.i18n import _
  12 from mercurial import cmdutil, commands, cmdutil, hg, mdiff, patch, revlog
  13 from mercurial import util, fancyopts, extensions
  14 import copy, cStringIO, errno, operator, os, re, shutil, tempfile
  15
  16 lines_re = re.compile(r'@@ -(\d+),(\d+) \+(\d+),(\d+) @@\s*(.*)')
  17
  18 def scanpatch(fp):
  19     lr = patch.linereader(fp)
  20
  21     def scanwhile(first, p):
  22         lines = [first]
  23         while True:
  24             line = lr.readline()
  25             if not line:
  26                 break
  27             if p(line):
  28                 lines.append(line)
  29             else:
  30                 lr.push(line)
  31                 break
  32         return lines
  33
  34     while True:
  35         line = lr.readline()
  36         if not line:
  37             break
  38         if line.startswith('diff --git a/'):
  39             def notheader(line):
  40                 s = line.split(None, 1)
  41                 return not s or s[0] not in ('---', 'diff')
  42             header = scanwhile(line, notheader)
  43             fromfile = lr.readline()
  44             if fromfile.startswith('---'):
  45                 tofile = lr.readline()
  46                 header += [fromfile, tofile]
  47             else:
  48                 lr.push(fromfile)
  49             yield 'file', header
  50         elif line[0] == ' ':
  51             yield 'context', scanwhile(line, lambda l: l[0] in ' \\')
  52         elif line[0] in '-+':
  53             yield 'hunk', scanwhile(line, lambda l: l[0] in '-+\\')
  54         else:
  55             m = lines_re.match(line)
  56             if m:
  57                 yield 'range', m.groups()
  58             else:
  59                 raise patch.PatchError('unknown patch content: %r' % line)
  60
  61 class header(object):
  62     diff_re = re.compile('diff --git a/(.*) b/(.*)$')
  63     allhunks_re = re.compile('(?:index|new file|deleted file) ')
  64     pretty_re = re.compile('(?:new file|deleted file) ')
  65     special_re = re.compile('(?:index|new|deleted|copy|rename) ')
  66
  67     def __init__(self, header):
  68         self.header = header
  69         self.hunks = []
  70
  71     def binary(self):
  72         for h in self.header:
  73             if h.startswith('index '):
  74                 return True
  75
  76     def pretty(self, fp):
  77         for h in self.header:
  78             if h.startswith('index '):
  79                 fp.write(_('this modifies a binary file (all or nothing)\n'))
  80                 break
  81             if self.pretty_re.match(h):
  82                 fp.write(h)
  83                 if self.binary():
  84                     fp.write(_('this is a binary file\n'))
  85                 break
  86             if h.startswith('---'):
  87                 fp.write(_('%d hunks, %d lines changed\n') %
  88                          (len(self.hunks),
  89                           sum([h.added + h.removed for h in self.hunks])))
  90                 break
  91             fp.write(h)
  92
  93     def write(self, fp):
  94         fp.write(''.join(self.header))
  95
  96     def allhunks(self):
  97         for h in self.header:
  98             if self.allhunks_re.match(h):
  99                 return True
 100
 101     def files(self):
 102         fromfile, tofile = self.diff_re.match(self.header[0]).groups()
 103         if fromfile == tofile:
 104             return [fromfile]
 105         return [fromfile, tofile]
 106
 107     def filename(self):
 108         return self.files()[-1]
 109
 110     def __repr__(self):
 111         return '<header %s>' % (' '.join(map(repr, self.files())))
 112
 113     def special(self):
 114         for h in self.header:
 115             if self.special_re.match(h):
 116                 return True
 117
 118 def countchanges(hunk):
 119     add = len([h for h in hunk if h[0] == '+'])
 120     rem = len([h for h in hunk if h[0] == '-'])
 121     return add, rem
 122
 123 class hunk(object):
 124     maxcontext = 3
 125
 126     def __init__(self, header, fromline, toline, proc, before, hunk, after):
 127         def trimcontext(number, lines):
 128             delta = len(lines) - self.maxcontext
 129             if False and delta > 0:
 130                 return number + delta, lines[:self.maxcontext]
 131             return number, lines
 132
 133         self.header = header
 134         self.fromline, self.before = trimcontext(fromline, before)
 135         self.toline, self.after = trimcontext(toline, after)
 136         self.proc = proc
 137         self.hunk = hunk
 138         self.added, self.removed = countchanges(self.hunk)
 139
 140     def __cmp__(self, rhs):
 141         # since the hunk().toline needs to be adjusted when hunks are
 142         # removed/added, we can't take it into account when we cmp
 143         attrs = ['header', 'fromline', 'proc', 'hunk', 'added', 'removed']
 144         for attr in attrs:
 145             selfattr = getattr(self, attr, None)
 146             rhsattr = getattr(rhs, attr, None)
 147
 148             if selfattr is None or rhsattr is None:
 149                 raise util.Abort(_('non-existant attribute %s') % attr)
 150
 151             rv = cmp(selfattr, rhsattr)
 152             if rv != 0:
 153                 return rv
 154         return rv
 155
 156
 157     def write(self, fp):
 158         delta = len(self.before) + len(self.after)
 159         if self.after and self.after[-1] == '\\ No newline at end of file\n':
 160             delta -= 1
 161         fromlen = delta + self.removed
 162         tolen = delta + self.added
 163         fp.write('@@ -%d,%d +%d,%d @@%s\n' %
 164                  (self.fromline, fromlen, self.toline, tolen,
 165                   self.proc and (' ' + self.proc)))
 166         fp.write(''.join(self.before + self.hunk + self.after))
 167
 168     pretty = write
 169
 170     def filename(self):
 171         return self.header.filename()
 172
 173     def __repr__(self):
 174         return '<hunk %r@%d>' % (self.filename(), self.fromline)
 175
 176 def parsepatch(fp):
 177     class parser(object):
 178         def __init__(self):
 179             self.fromline = 0
 180             self.toline = 0
 181             self.proc = ''
 182             self.header = None
 183             self.context = []
 184             self.before = []
 185             self.hunk = []
 186             self.stream = []
 187
 188         def addrange(self, (fromstart, fromend, tostart, toend, proc)):
 189             self.fromline = int(fromstart)
 190             self.toline = int(tostart)
 191             self.proc = proc
 192
 193         def addcontext(self, context):
 194             if self.hunk:
 195                 h = hunk(self.header, self.fromline, self.toline, self.proc,
 196                          self.before, self.hunk, context)
 197                 self.header.hunks.append(h)
 198                 self.stream.append(h)
 199                 self.fromline += len(self.before) + h.removed
 200                 self.toline += len(self.before) + h.added
 201                 self.before = []
 202                 self.hunk = []
 203                 self.proc = ''
 204             self.context = context
 205
 206         def addhunk(self, hunk):
 207             if self.context:
 208                 self.before = self.context
 209                 self.context = []
 210             self.hunk = hunk
 211
 212         def newfile(self, hdr):
 213             self.addcontext([])
 214             h = header(hdr)
 215             self.stream.append(h)
 216             self.header = h
 217
 218         def finished(self):
 219             self.addcontext([])
 220             return self.stream
 221
 222         transitions = {
 223             'file': {'context': addcontext,
 224                      'file': newfile,
 225                      'hunk': addhunk,
 226                      'range': addrange},
 227             'context': {'file': newfile,
 228                         'hunk': addhunk,
 229                         'range': addrange},
 230             'hunk': {'context': addcontext,
 231                      'file': newfile,
 232                      'range': addrange},
 233             'range': {'context': addcontext,
 234                       'hunk': addhunk},
 235             }
 236
 237     p = parser()
 238
 239     state = 'context'
 240     for newstate, data in scanpatch(fp):
 241         try:
 242             p.transitions[state][newstate](p, data)
 243         except KeyError:
 244             raise patch.PatchError('unhandled transition: %s -> %s' %
 245                                    (state, newstate))
 246         state = newstate
 247     return p.finished()
 248
 249 def filterpatch(ui, chunks, shouldprompt=True):
 250     chunks = list(chunks)
 251     chunks.reverse()
 252     seen = {}
 253     def consumefile():
 254         consumed = []
 255         while chunks:
 256             if isinstance(chunks[-1], header):
 257                 break
 258             else:
 259                 consumed.append(chunks.pop())
 260         return consumed
 261
 262     resp_all = [None]
 263
 264     """ If we're not to prompt (i.e. they specified the --all flag)
 265         we pre-emptively set the 'all' flag """
 266     if shouldprompt == False:
 267         resp_all = ['y']
 268
 269     resp_file = [None]
 270     applied = {}
 271     def prompt(query):
 272         if resp_all[0] is not None:
 273             return resp_all[0]
 274         if resp_file[0] is not None:
 275             return resp_file[0]
 276         while True:
 277             resps = _('[Ynsfdaq?]')
 278             choices = (_('&Yes, shelve this change'),
 279                     _('&No, skip this change'),
 280                     _('&Skip remaining changes to this file'),
 281                     _('Shelve remaining changes to this &file'),
 282                     _('&Done, skip remaining changes and files'),
 283                     _('Shelve &all changes to all remaining files'),
 284                     _('&Quit, shelving no changes'),
 285                     _('&?'))
 286             r = ui.promptchoice("%s %s " % (query, resps), choices)
 287             if r == 7:
 288                 c = shelve.__doc__.find('y - shelve this change')
 289                 for l in shelve.__doc__[c:].splitlines():
 290                     if l: ui.write(_(l.strip()) + '\n')
 291                 continue
 292             elif r == 0: # yes
 293                 ret = 'y'
 294             elif r == 1: # no
 295                 ret = 'n'
 296             elif r == 2: # Skip
 297                 ret = resp_file[0] = 'n'
 298             elif r == 3: # file (shelve remaining)
 299                 ret = resp_file[0] = 'y'
 300             elif r == 4: # done, skip remaining
 301                 ret = resp_all[0] = 'n'
 302             elif r == 5: # all
 303                 ret = resp_all[0] = 'y'
 304             elif r == 6: # quit
 305                 raise util.Abort(_('user quit'))
 306             return ret
 307     while chunks:
 308         chunk = chunks.pop()
 309         if isinstance(chunk, header):
 310             resp_file = [None]
 311             fixoffset = 0
 312             hdr = ''.join(chunk.header)
 313             if hdr in seen:
 314                 consumefile()
 315                 continue
 316             seen[hdr] = True
 317             if resp_all[0] is None:
 318                 chunk.pretty(ui)
 319             if shouldprompt == True:
 320                 r = prompt(_('shelve changes to %s?') %
 321                        _(' and ').join(map(repr, chunk.files())))
 322             else:
 323                 r = 'y'
 324
 325             if r == 'y':
 326                 applied[chunk.filename()] = [chunk]
 327                 if chunk.allhunks():
 328                     applied[chunk.filename()] += consumefile()
 329             else:
 330                 consumefile()
 331         else:
 332             if resp_file[0] is None and resp_all[0] is None:
 333                 chunk.pretty(ui)
 334             r = prompt(_('shelve this change to %r?') %
 335                        chunk.filename())
 336             if r == 'y':
 337                 if fixoffset:
 338                     chunk = copy.copy(chunk)
 339                     chunk.toline += fixoffset
 340                 applied[chunk.filename()].append(chunk)
 341             else:
 342                 fixoffset += chunk.removed - chunk.added
 343     return reduce(operator.add, [h for h in applied.itervalues()
 344                                  if h[0].special() or len(h) > 1], [])
 345
 346 def refilterpatch(allchunk, selected):
 347     ''' return unshelved chunks of files to be shelved '''
 348     l = []
 349     fil = []
 350     for c in allchunk:
 351         if isinstance(c, header):
 352             if len(l) > 1 and l[0] in selected:
 353                 fil += l
 354             l = [c]
 355         elif c not in selected:
 356             l.append(c)
 357     if len(l) > 1 and l[0] in selected:
 358         fil += l
 359     return fil
 360
 361 def makebackup(ui, repo, dir, files):
 362     try:
 363         os.mkdir(dir)
 364     except OSError, err:
 365         if err.errno != errno.EEXIST:
 366             raise
 367
 368     backups = {}
 369     for f in files:
 370         fd, tmpname = tempfile.mkstemp(prefix=f.replace('/', '_')+'.',
 371                                        dir=dir)
 372         os.close(fd)
 373         ui.debug('backup %r as %r\n' % (f, tmpname))
 374         util.copyfile(repo.wjoin(f), tmpname)
 375         backups[f] = tmpname
 376
 377     return backups
 378
 379 def getshelfpath(repo, name):
 380     if name:
 381         shelfpath = "shelves/" + name
 382     else:
 383         # Check if a shelf from an older version exists
 384         if os.path.isfile(repo.join('shelve')):
 385             shelfpath = 'shelve'
 386         else:
 387             shelfpath = "shelves/default"
 388
 389     return shelfpath
 390
 391 def shelve(ui, repo, *pats, **opts):
 392     '''interactively select changes to set aside
 393
 394     If a list of files is omitted, all changes reported by "hg status"
 395     will be candidates for shelving.
 396
 397     You will be prompted for whether to shelve changes to each
 398     modified file, and for files with multiple changes, for each
 399     change to use.
 400
 401     The shelve command works with the Color extension to display
 402     diffs in color.
 403
 404     On each prompt, the following responses are possible::
 405
 406       y - shelve this change
 407       n - skip this change
 408
 409       s - skip remaining changes to this file
 410       f - shelve remaining changes to this file
 411
 412       d - done, skip remaining changes and files
 413       a - shelve all changes to all remaining files
 414       q - quit, shelving no changes
 415
 416       ? - display help'''
 417
 418     if not ui.interactive:
 419         raise util.Abort(_('shelve can only be run interactively'))
 420
 421     # List all the active shelves by name and return '
 422     if opts['list']:
 423         listshelves(ui,repo)
 424         return
 425
 426     forced = opts['force'] or opts['append']
 427
 428     # Shelf name and path
 429     shelfname = opts.get('name')
 430     shelfpath = getshelfpath(repo, shelfname)
 431
 432     if os.path.exists(repo.join(shelfpath)) and not forced:
 433         raise util.Abort(_('shelve data already exists'))
 434
 435     def shelvefunc(ui, repo, message, match, opts):
 436         changes = repo.status(match=match)[:5]
 437         modified, added, removed = changes[:3]
 438         files = modified + added + removed
 439         diffopts = mdiff.diffopts(git=True, nodates=True)
 440         patch_diff = ''.join(patch.diff(repo, repo.dirstate.parents()[0],
 441                            match=match, changes=changes, opts=diffopts))
 442
 443         fp = cStringIO.StringIO(patch_diff)
 444         ac = parsepatch(fp)
 445         fp.close()
 446
 447         chunks = filterpatch(ui, ac, not opts['all'])
 448         rc = refilterpatch(ac, chunks)
 449
 450         contenders = {}
 451         for h in chunks:
 452             try: contenders.update(dict.fromkeys(h.files()))
 453             except AttributeError: pass
 454
 455         newfiles = [f for f in files if f in contenders]
 456
 457         if not newfiles:
 458             ui.status(_('no changes to shelve\n'))
 459             return 0
 460
 461         modified = dict.fromkeys(changes[0])
 462
 463         backupdir = repo.join('shelve-backups')
 464
 465         try:
 466             bkfiles = [f for f in newfiles if f in modified]
 467             backups = makebackup(ui, repo, backupdir, bkfiles)
 468
 469             # patch to shelve
 470             sp = cStringIO.StringIO()
 471             for c in chunks:
 472                 if c.filename() in backups:
 473                     c.write(sp)
 474             doshelve = sp.tell()
 475             sp.seek(0)
 476
 477             # patch to apply to shelved files
 478             fp = cStringIO.StringIO()
 479             for c in rc:
 480                 if c.filename() in backups:
 481                     c.write(fp)
 482             dopatch = fp.tell()
 483             fp.seek(0)
 484
 485             try:
 486                 # 3a. apply filtered patch to clean repo (clean)
 487                 if backups:
 488                     hg.revert(repo, repo.dirstate.parents()[0], backups.has_key)
 489
 490                 # 3b. apply filtered patch to clean repo (apply)
 491                 if dopatch:
 492                     ui.debug('applying patch\n')
 493                     ui.debug(fp.getvalue())
 494                     patch.internalpatch(ui, repo, fp, 1)
 495                 del fp
 496
 497                 # 3c. apply filtered patch to clean repo (shelve)
 498                 if doshelve:
 499                     ui.debug("saving patch to shelve\n")
 500                     if opts['append']:
 501                         f = repo.opener(shelfpath, "a")
 502                     else:
 503                         f = repo.opener(shelfpath, "w")
 504                     f.write(sp.getvalue())
 505                     del f
 506                 del sp
 507             except:
 508                 try:
 509                     for realname, tmpname in backups.iteritems():
 510                         ui.debug('restoring %r to %r\n' % (tmpname, realname))
 511                         util.copyfile(tmpname, repo.wjoin(realname))
 512                     ui.debug('removing shelve file\n')
 513                     os.unlink(repo.join(shelfpath))
 514                 except OSError:
 515                     pass
 516
 517             return 0
 518         finally:
 519             try:
 520                 for realname, tmpname in backups.iteritems():
 521                     ui.debug('removing backup for %r : %r\n' % (realname, tmpname))
 522                     os.unlink(tmpname)
 523                 os.rmdir(backupdir)
 524             except OSError:
 525                 pass
 526     fancyopts.fancyopts([], commands.commitopts, opts)
 527
 528     # wrap ui.write so diff output can be labeled/colorized
 529     def wrapwrite(orig, *args, **kw):
 530         label = kw.pop('label', '')
 531         if label: label += ' '
 532         for chunk, l in patch.difflabel(lambda: args):
 533             orig(chunk, label=label + l)
 534     oldwrite = ui.write
 535     extensions.wrapfunction(ui, 'write', wrapwrite)
 536     try:
 537         return cmdutil.commit(ui, repo, shelvefunc, pats, opts)
 538     finally:
 539         ui.write = oldwrite
 540
 541 def listshelves(ui, repo):
 542     # Check for shelve file at old location first
 543     if os.path.isfile(repo.join('shelve')):
 544         ui.status('default\n')
 545
 546     # Now go through all the files in the shelves folder and list them out
 547     dirname = repo.join('shelves')
 548     if os.path.isdir(dirname):
 549         for filename in os.listdir(repo.join('shelves')):
 550             ui.status(filename + '\n')
 551
 552 def unshelve(ui, repo, **opts):
 553     '''restore shelved changes'''
 554
 555     # Shelf name and path
 556     shelfname = opts.get('name')
 557     shelfpath = getshelfpath(repo, shelfname)
 558
 559     # List all the active shelves by name and return '
 560     if opts['list']:
 561         listshelves(ui,repo)
 562         return
 563
 564     try:
 565         patch_diff = repo.opener(shelfpath).read()
 566         fp = cStringIO.StringIO(patch_diff)
 567         if opts['inspect']:
 568             # wrap ui.write so diff output can be labeled/colorized
 569             def wrapwrite(orig, *args, **kw):
 570                 label = kw.pop('label', '')
 571                 if label: label += ' '
 572                 for chunk, l in patch.difflabel(lambda: args):
 573                     orig(chunk, label=label + l)
 574             oldwrite = ui.write
 575             extensions.wrapfunction(ui, 'write', wrapwrite)
 576             try:
 577                 ui.status(fp.getvalue())
 578             finally:
 579                 ui.write = oldwrite
 580         else:
 581             files = []
 582             ac = parsepatch(fp)
 583             for chunk in ac:
 584                 if isinstance(chunk, header):
 585                     files += chunk.files()
 586             backupdir = repo.join('shelve-backups')
 587             backups = makebackup(ui, repo, backupdir, set(files))
 588
 589             ui.debug('applying shelved patch\n')
 590             patchdone = 0
 591             try:
 592                 try:
 593                     fp.seek(0)
 594                     patch.internalpatch(ui, repo, fp, 1)
 595                     patchdone = 1
 596                 except:
 597                     if opts['force']:
 598                         patchdone = 1
 599                     else:
 600                         ui.status('restoring backup files\n')
 601                         for realname, tmpname in backups.iteritems():
 602                             ui.debug('restoring %r to %r\n' %
 603                                      (tmpname, realname))
 604                             util.copyfile(tmpname, repo.wjoin(realname))
 605             finally:
 606                 try:
 607                     ui.debug('removing backup files\n')
 608                     shutil.rmtree(backupdir, True)
 609                 except OSError:
 610                     pass
 611
 612             if patchdone:
 613                 ui.debug("removing shelved patches\n")
 614                 os.unlink(repo.join(shelfpath))
 615                 ui.status("unshelve completed\n")
 616     except IOError:
 617         ui.warn('nothing to unshelve\n')
 618
 619 cmdtable = {
 620     "shelve":
 621         (shelve,
 622          [('A', 'addremove', None,
 623            _('mark new/missing files as added/removed before shelving')),
 624           ('f', 'force', None,
 625            _('overwrite existing shelve data')),
 626           ('a', 'append', None,
 627            _('append to existing shelve data')),
 628           ('', 'all', None,
 629            _('shelve all changes')),
 630           ('n', 'name', '',
 631            _('shelve changes to specified shelf name')),
 632            ('l', 'list', None, _('list active shelves')),
 633          ] + commands.walkopts,
 634          _('hg shelve [OPTION]... [FILE]...')),
 635     "unshelve":
 636         (unshelve,
 637          [('i', 'inspect', None, _('inspect shelved changes only')),
 638           ('f', 'force', None,
 639            _('proceed even if patches do not unshelve cleanly')),
 640            ('n', 'name', '',
 641             _('unshelve changes from specified shelf name')),
 642            ('l', 'list', None, _('list active shelves')),
 643          ],
 644          _('hg unshelve [OPTION]...')),
 645 }