shelve: preserve newly created branch on non-bare shelve in wctx (BC)

shelve: preserve newly created branch on non-bare shelve in wctx (BC)

Before this patch current branch in working context wasnt preserved
after shelve, this patch makes it restore after update.

--- a/hgext/shelve.py	Fri Mar 18 16:14:57 2016 -0400
+++ b/hgext/shelve.py	Sat Mar 12 18:36:31 2016 +0100
@@ -272,6 +272,7 @@
     if len(parents) > 1:
         raise error.Abort(_('cannot shelve while merging'))
     parent = parents[0]
+    origbranch = wctx.branch()
 
     # we never need the user, so we use a generic user for all shelve operations
     user = 'shelve@localhost'
@@ -378,11 +379,19 @@
             desc = util.ellipsis(desc, ui.termwidth())
         ui.status(_('shelved as %s\n') % name)
         hg.update(repo, parent.node())
+        if origbranch != repo['.'].branch() and not _isbareshelve(pats, opts):
+            repo.dirstate.setbranch(origbranch)
 
         _aborttransaction(repo)
     finally:
         lockmod.release(tr, lock)
 
+def _isbareshelve(pats, opts):
+    return (not pats
+            and not opts.get('interactive', False)
+            and not opts.get('include', False)
+            and not opts.get('exclude', False))
+
 def cleanupcmd(ui, repo):
     """subcommand that deletes all shelves"""
 

--- a/tests/test-shelve.t	Fri Mar 18 16:14:57 2016 -0400
+++ b/tests/test-shelve.t	Sat Mar 12 18:36:31 2016 +0100
@@ -1318,3 +1318,62 @@
   $ hg commit -qm "Remove unknown"
 
   $ cd ..
+
+We expects that non-bare shelve keeps newly created branch in
+working directory.
+
+  $ hg init shelve-preserve-new-branch
+  $ cd shelve-preserve-new-branch
+  $ echo "a" >> a
+  $ hg add a
+  $ echo "b" >> b
+  $ hg add b
+  $ hg commit -m "ab"
+  $ echo "aa" >> a
+  $ echo "bb" >> b
+  $ hg branch new-branch
+  marked working directory as branch new-branch
+  (branches are permanent and global, did you want a bookmark?)
+  $ hg status
+  M a
+  M b
+  $ hg branch
+  new-branch
+  $ hg shelve a
+  shelved as default
+  1 files updated, 0 files merged, 0 files removed, 0 files unresolved
+  $ hg branch
+  new-branch
+  $ hg status
+  M b
+  $ touch "c" >> c
+  $ hg add c
+  $ hg status
+  M b
+  A c
+  $ hg shelve --exclude c
+  shelved as default-01
+  1 files updated, 0 files merged, 0 files removed, 0 files unresolved
+  $ hg branch
+  new-branch
+  $ hg status
+  A c
+  $ hg shelve --include c
+  shelved as default-02
+  0 files updated, 0 files merged, 1 files removed, 0 files unresolved
+  $ hg branch
+  new-branch
+  $ hg status
+  $ echo "d" >> d
+  $ hg add d
+  $ hg status
+  A d
+
+We expect that bare-shelve will not keep branch in current working directory.
+
+  $ hg shelve
+  shelved as default-03
+  0 files updated, 0 files merged, 1 files removed, 0 files unresolved
+  $ hg branch
+  default
+