Add gitbranch option

README.md

@@ -101,6 +101,7 @@ optional arguments:
 -m,  --machine         prints system and machine info
 -g,  --githash         prints current Git commit hash
 -r,  --gitrepo         prints current Git remote address
+-b,  --gitbranch       prints the current Git branch (new in v1.6)
 -iv, --iversion        print name and version of all imported packages      
 -w,  --watermark       prints the current version of watermark
 ```

watermark/__init__.py

@@ -9,7 +9,7 @@
 import sys
 
 
-__version__ = '1.5.0'
+__version__ = '1.6.0dev'
 
 if sys.version_info >= (3, 0):
     from watermark.watermark import *

watermark/watermark.py

@@ -67,6 +67,8 @@ class WaterMark(Magics):
               help='prints current Git commit hash')
     @argument('-r', '--gitrepo', action='store_true',
               help='prints current Git remote address')
+    @argument('-b', '--gitbranch', action='store_true',
+              help='prints current Git branch')
     @argument('-w', '--watermark', action='store_true',
               help='prints the current version of watermark')
     @argument('-iv', '--iversions', action='store_true',
@@ -127,6 +129,8 @@ def watermark(self, line):
                 self._get_commit_hash(bool(args.machine))
             if args.gitrepo:
                 self._get_git_remote_origin(bool(args.machine))
+            if args.gitbranch:
+                self._get_git_branch(bool(args.machine))
             if args.iversions:
                 self._print_all_import_versions(self.shell.user_ns)
             if args.watermark:
@@ -209,6 +213,18 @@ def _get_git_remote_origin(self, machine):
         self.out += '\nGit repo%s: %s' % (space,
                                           git_remote_origin.decode("utf-8"))
 
+    def _get_git_branch(self, machine):
+        process = subprocess.Popen(['git', 'rev-parse', '--abbrev-ref',
+                                    'HEAD'],
+                                   shell=False,
+                                   stdout=subprocess.PIPE)
+        git_branch = process.communicate()[0].strip()
+        space = ''
+        if machine:
+            space = ' '
+        self.out += '\nGit branch%s: %s' % (space,
+                                            git_branch.decode("utf-8"))
+
     @staticmethod
     def _print_all_import_versions(vars):
         for val in list(vars.values()):