--- a/hgext/rebase.py	Wed Dec 03 22:03:05 2008 -0800
+++ b/hgext/rebase.py	Fri Dec 05 22:22:14 2008 +0100
@@ -64,6 +64,14 @@
         contf = opts.get('continue')
         abortf = opts.get('abort')
         collapsef = opts.get('collapse', False)
+        extrafn = opts.get('extrafn')
+        if opts.get('keepbranches', None):
+            if extrafn:
+                raise dispatch.ParseError('rebase',
+                        _('cannot use both keepbranches and extrafn'))
+            def extrafn(ctx, extra):
+                extra['branch'] = ctx.branch()
+
         if contf or abortf:
             if contf and abortf:
                 raise dispatch.ParseError('rebase',
@@ -101,14 +109,14 @@
                 storestatus(repo, originalwd, target, state, collapsef,
                                                                 external)
                 rebasenode(repo, rev, target, state, skipped, targetancestors,
-                                                                collapsef)
+                                                       collapsef, extrafn)
         ui.note(_('rebase merging completed\n'))
 
         if collapsef:
             p1, p2 = defineparents(repo, min(state), target,
                                                         state, targetancestors)
             concludenode(repo, rev, p1, external, state, collapsef,
-                                                last=True, skipped=skipped)
+                         last=True, skipped=skipped, extrafn=extrafn)
 
         if 'qtip' in repo.tags():
             updatemq(repo, state, skipped, **opts)
@@ -131,7 +139,8 @@
     finally:
         del lock, wlock
 
-def concludenode(repo, rev, p1, p2, state, collapse, last=False, skipped={}):
+def concludenode(repo, rev, p1, p2, state, collapse, last=False, skipped={},
+                 extrafn=None):
     """Skip commit if collapsing has been required and rev is not the last
     revision, commit otherwise
     """
@@ -155,18 +164,22 @@
         else:
             commitmsg = repo[rev].description()
         # Commit might fail if unresolved files exist
+        extra = {'rebase_source': repo[rev].hex()}
+        if extrafn:
+            extrafn(repo[rev], extra)
         newrev = repo.commit(m+a+r,
                             text=commitmsg,
                             user=repo[rev].user(),
                             date=repo[rev].date(),
-                            extra={'rebase_source': repo[rev].hex()})
+                            extra=extra)
         return newrev
     except util.Abort:
         # Invalidate the previous setparents
         repo.dirstate.invalidate()
         raise
 
-def rebasenode(repo, rev, target, state, skipped, targetancestors, collapse):
+def rebasenode(repo, rev, target, state, skipped, targetancestors, collapse,
+               extrafn):
     'Rebase a single revision'
     repo.ui.debug(_("rebasing %d:%s\n") % (rev, repo[rev]))
 
@@ -195,7 +208,8 @@
         repo.ui.debug(_('resuming interrupted rebase\n'))
 
 
-    newrev = concludenode(repo, rev, p1, p2, state, collapse)
+    newrev = concludenode(repo, rev, p1, p2, state, collapse,
+                          extrafn=extrafn)
 
     # Update the state
     if newrev is not None:
@@ -409,6 +423,7 @@
         (rebase,
         [
         ('', 'keep', False, _('keep original revisions')),
+        ('', 'keepbranches', False, _('keep original branches')),
         ('s', 'source', '', _('rebase from a given revision')),
         ('b', 'base', '', _('rebase from the base of a given revision')),
         ('d', 'dest', '', _('rebase onto a given revision')),

--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/tests/test-rebase-keep-branch	Fri Dec 05 22:22:14 2008 +0100
@@ -0,0 +1,30 @@
+#!/bin/sh
+
+echo "[extensions]" >> $HGRCPATH
+echo "graphlog=" >> $HGRCPATH
+echo "rebase=" >> $HGRCPATH
+
+addcommit () {
+    echo $1 > $1
+    hg add $1
+    hg commit -d "${2} 0" -u test -m $1
+}
+
+hg init a
+cd a
+addcommit "c1" 0
+addcommit "c2" 1
+
+addcommit "l1" 2
+addcommit "l2" 3
+
+hg update -C 1
+hg branch 'notdefault'
+addcommit "r1" 4
+hg glog --template '{rev}:{desc}:{branches}\n'
+
+echo
+echo '% Rebase a branch while preserving the branch name'
+hg update -C 3
+hg rebase -b 4 -d 3 --keepbranches 2>&1 | sed 's/\(saving bundle to \).*/\1/'
+hg glog --template '{rev}:{desc}:{branches}\n'

--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/tests/test-rebase-keep-branch.out	Fri Dec 05 22:22:14 2008 +0100
@@ -0,0 +1,33 @@
+0 files updated, 0 files merged, 2 files removed, 0 files unresolved
+marked working directory as branch notdefault
+created new head
+@  4:r1:notdefault
+|
+| o  3:l2:
+| |
+| o  2:l1:
+|/
+o  1:c2:
+|
+o  0:c1:
+
+
+% Rebase a branch while preserving the branch name
+2 files updated, 0 files merged, 1 files removed, 0 files unresolved
+saving bundle to 
+adding branch
+adding changesets
+adding manifests
+adding file changes
+added 1 changesets with 1 changes to 1 files
+rebase completed
+@  4:r1:notdefault
+|
+o  3:l2:
+|
+o  2:l1:
+|
+o  1:c2:
+|
+o  0:c1:
+

--- a/tests/test-rebase-parameters.out	Wed Dec 03 22:03:05 2008 -0800
+++ b/tests/test-rebase-parameters.out	Fri Dec 05 22:22:14 2008 +0100
@@ -15,15 +15,16 @@
 
 options:
 
-    --keep      keep original revisions
- -s --source    rebase from a given revision
- -b --base      rebase from the base of a given revision
- -d --dest      rebase onto a given revision
-    --collapse  collapse the rebased revisions
- -c --continue  continue an interrupted rebase
- -a --abort     abort an interrupted rebase
-    --style     display using template map file
-    --template  display with template
+    --keep          keep original revisions
+    --keepbranches  keep original branches
+ -s --source        rebase from a given revision
+ -b --base          rebase from the base of a given revision
+ -d --dest          rebase onto a given revision
+    --collapse      collapse the rebased revisions
+ -c --continue      continue an interrupted rebase
+ -a --abort         abort an interrupted rebase
+    --style         display using template map file
+    --template      display with template
 
 use "hg -v help rebase" to show global options
 
@@ -42,15 +43,16 @@
 
 options:
 
-    --keep      keep original revisions
- -s --source    rebase from a given revision
- -b --base      rebase from the base of a given revision
- -d --dest      rebase onto a given revision
-    --collapse  collapse the rebased revisions
- -c --continue  continue an interrupted rebase
- -a --abort     abort an interrupted rebase
-    --style     display using template map file
-    --template  display with template
+    --keep          keep original revisions
+    --keepbranches  keep original branches
+ -s --source        rebase from a given revision
+ -b --base          rebase from the base of a given revision
+ -d --dest          rebase onto a given revision
+    --collapse      collapse the rebased revisions
+ -c --continue      continue an interrupted rebase
+ -a --abort         abort an interrupted rebase
+    --style         display using template map file
+    --template      display with template
 
 use "hg -v help rebase" to show global options
 
@@ -69,15 +71,16 @@
 
 options:
 
-    --keep      keep original revisions
- -s --source    rebase from a given revision
- -b --base      rebase from the base of a given revision
- -d --dest      rebase onto a given revision
-    --collapse  collapse the rebased revisions
- -c --continue  continue an interrupted rebase
- -a --abort     abort an interrupted rebase
-    --style     display using template map file
-    --template  display with template
+    --keep          keep original revisions
+    --keepbranches  keep original branches
+ -s --source        rebase from a given revision
+ -b --base          rebase from the base of a given revision
+ -d --dest          rebase onto a given revision
+    --collapse      collapse the rebased revisions
+ -c --continue      continue an interrupted rebase
+ -a --abort         abort an interrupted rebase
+    --style         display using template map file
+    --template      display with template
 
 use "hg -v help rebase" to show global options
 
@@ -96,15 +99,16 @@
 
 options:
 
-    --keep      keep original revisions
- -s --source    rebase from a given revision
- -b --base      rebase from the base of a given revision
- -d --dest      rebase onto a given revision
-    --collapse  collapse the rebased revisions
- -c --continue  continue an interrupted rebase
- -a --abort     abort an interrupted rebase
-    --style     display using template map file
-    --template  display with template
+    --keep          keep original revisions
+    --keepbranches  keep original branches
+ -s --source        rebase from a given revision
+ -b --base          rebase from the base of a given revision
+ -d --dest          rebase onto a given revision
+    --collapse      collapse the rebased revisions
+ -c --continue      continue an interrupted rebase
+ -a --abort         abort an interrupted rebase
+    --style         display using template map file
+    --template      display with template
 
 use "hg -v help rebase" to show global options