From 0e20aa09e25b06da652702b3bbb6260921bd0d4d Mon Sep 17 00:00:00 2001
From: Oscar Landry <lasm.landry@gmail.com>
Date: Wed, 26 Apr 2017 10:04:26 +0200
Subject: [PATCH] define more branch protected

---
 pre_commit_hooks/no_commit_to_branch.py | 18 ++++++++++++------
 tests/check_no_commit_to_branch_test.py |  8 ++++----
 2 files changed, 16 insertions(+), 10 deletions(-)

diff --git a/pre_commit_hooks/no_commit_to_branch.py b/pre_commit_hooks/no_commit_to_branch.py
index 4aa3535..b77d0b6 100644
--- a/pre_commit_hooks/no_commit_to_branch.py
+++ b/pre_commit_hooks/no_commit_to_branch.py
@@ -6,20 +6,26 @@ import sys
 from pre_commit_hooks.util import cmd_output
 
 
-def is_on_branch(protected):
+def is_on_branch(protected=()):
     branch = cmd_output('git', 'symbolic-ref', 'HEAD')
     chunks = branch.strip().split('/')
-    return '/'.join(chunks[2:]) == protected
+    position = '/'.join(chunks[2:])
+    return position in (protected or ('master',))
 
 
-def main(argv=[]):
+def main(argv=None):
     parser = argparse.ArgumentParser()
     parser.add_argument(
-        '-b', '--branch', default='master', help='branch to disallow commits to')
+        '-b',
+        '--branch',
+        action='append',
+        dest='branches',
+        help='branch to disallow commits to'
+    )
     args = parser.parse_args(argv)
 
-    return int(is_on_branch(args.branch))
+    return int(is_on_branch(args.branches))
 
 
 if __name__ == '__main__':
-    sys.exit(main(sys.argv))
+    sys.exit(main(sys.argv[1:]))
diff --git a/tests/check_no_commit_to_branch_test.py b/tests/check_no_commit_to_branch_test.py
index 99af938..d33a8be 100644
--- a/tests/check_no_commit_to_branch_test.py
+++ b/tests/check_no_commit_to_branch_test.py
@@ -9,24 +9,24 @@ from pre_commit_hooks.util import cmd_output
 def test_other_branch(temp_git_dir):
     with temp_git_dir.as_cwd():
         cmd_output('git', 'checkout', '-b', 'anotherbranch')
-        assert is_on_branch('master') is False
+        assert is_on_branch(('master', )) is False
 
 
 def test_multi_branch(temp_git_dir):
     with temp_git_dir.as_cwd():
         cmd_output('git', 'checkout', '-b', 'another/branch')
-        assert is_on_branch('master') is False
+        assert is_on_branch(('master', )) is False
 
 
 def test_multi_branch_fail(temp_git_dir):
     with temp_git_dir.as_cwd():
         cmd_output('git', 'checkout', '-b', 'another/branch')
-        assert is_on_branch('another/branch') is True
+        assert is_on_branch(('another/branch', )) is True
 
 
 def test_master_branch(temp_git_dir):
     with temp_git_dir.as_cwd():
-        assert is_on_branch('master') is True
+        assert is_on_branch(('master', )) is True
 
 
 def test_main_b_call(temp_git_dir):