Allow multiple branches to be protected

Original patch by @moas

README.md

@@ -82,7 +82,10 @@ Add this to your `.pre-commit-config.yaml`
 - `name-tests-test` - Assert that files in tests/ end in `_test.py`.
     - Use `args: ['--django']` to match `test*.py` instead.
 - `no-commit-to-branch` - Protect specific branches from direct checkins.
-    - Use `args: -b <branch> ` to set the branch. `master` is the default if no argument is set.
+    - Use `args: [--branch <branch>]` to set the branch. `master` is the
+      default if no argument is set.
+    - `-b` / `--branch` may be specified multiple times to protect multiple
+      branches.
 - `pyflakes` - Run pyflakes on your python files.
 - `pretty-format-json` - Checks that all your JSON files are pretty.  "Pretty"
   here means that keys are sorted and indented.  You can configure this with

pre_commit_hooks/no_commit_to_branch.py

@@ -12,18 +12,19 @@ def is_on_branch(protected):
     except CalledProcessError:
         return False
     chunks = branch.strip().split('/')
-    return '/'.join(chunks[2:]) == protected
+    return '/'.join(chunks[2:]) in protected
 
 
 def main(argv=None):
     parser = argparse.ArgumentParser()
     parser.add_argument(
-        '-b', '--branch', default='master',
+        '-b', '--branch', action='append',
-        help='branch to disallow commits to',
+        help='branch to disallow commits to, may be specified multiple times',
     )
     args = parser.parse_args(argv)
 
-    return int(is_on_branch(args.branch))
+    protected = set(args.branch or ('master',))
+    return int(is_on_branch(protected))
 
 
 if __name__ == '__main__':

tests/no_commit_to_branch_test.py

@@ -1,6 +1,8 @@
 from __future__ import absolute_import
 from __future__ import unicode_literals
 
+import pytest
+
 from pre_commit_hooks.no_commit_to_branch import is_on_branch
 from pre_commit_hooks.no_commit_to_branch import main
 from pre_commit_hooks.util import cmd_output
@@ -9,24 +11,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_branch_call(temp_git_dir):
@@ -35,6 +37,13 @@ def test_main_branch_call(temp_git_dir):
         assert main(('--branch', 'other')) == 1
 
 
+@pytest.mark.parametrize('branch_name', ('b1', 'b2'))
+def test_forbid_multiple_branches(temp_git_dir, branch_name):
+    with temp_git_dir.as_cwd():
+        cmd_output('git', 'checkout', '-b', branch_name)
+        assert main(('--branch', 'b1', '--branch', 'b2'))
+
+
 def test_main_default_call(temp_git_dir):
     with temp_git_dir.as_cwd():
         cmd_output('git', 'checkout', '-b', 'anotherbranch')