Merge pull request #376 from marcjay/no-commit-to-branch-wildcard

Add regex matching to no-commit-to-branch hook
diff --git a/README.md b/README.md
index a937771..5bf28d3 100644
--- a/README.md
+++ b/README.md
@@ -80,9 +80,11 @@
     - Use `args: ['--django']` to match `test*.py` instead.
 - `no-commit-to-branch` - Protect specific branches from direct checkins.
     - Use `args: [--branch, staging, --branch, master]` to set the branch.
-      `master` is the default if no argument is set.
+      `master` is the default if no branch argument is set.
     - `-b` / `--branch` may be specified multiple times to protect multiple
       branches.
+    - `-p` / `--pattern` can be used to protect branches that match a supplied regex
+      (e.g. `--pattern, release/.*`). May be specified multiple times.
 - `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
   the following commandline options:
diff --git a/pre_commit_hooks/no_commit_to_branch.py b/pre_commit_hooks/no_commit_to_branch.py
index 6b68c91..3131059 100644
--- a/pre_commit_hooks/no_commit_to_branch.py
+++ b/pre_commit_hooks/no_commit_to_branch.py
@@ -1,21 +1,26 @@
 from __future__ import print_function
 
 import argparse
+import re
+from typing import AbstractSet
 from typing import Optional
 from typing import Sequence
-from typing import Set
 
 from pre_commit_hooks.util import CalledProcessError
 from pre_commit_hooks.util import cmd_output
 
 
-def is_on_branch(protected):  # type: (Set[str]) -> bool
+def is_on_branch(protected, patterns=frozenset()):
+    # type: (AbstractSet[str], AbstractSet[str]) -> bool
     try:
-        branch = cmd_output('git', 'symbolic-ref', 'HEAD')
+        ref_name = cmd_output('git', 'symbolic-ref', 'HEAD')
     except CalledProcessError:
         return False
-    chunks = branch.strip().split('/')
-    return '/'.join(chunks[2:]) in protected
+    chunks = ref_name.strip().split('/')
+    branch_name = '/'.join(chunks[2:])
+    return branch_name in protected or any(
+        re.match(p, branch_name) for p in patterns
+    )
 
 
 def main(argv=None):  # type: (Optional[Sequence[str]]) -> int
@@ -24,10 +29,18 @@
         '-b', '--branch', action='append',
         help='branch to disallow commits to, may be specified multiple times',
     )
+    parser.add_argument(
+        '-p', '--pattern', action='append',
+        help=(
+            'regex pattern for branch name to disallow commits to, '
+            'may be specified multiple times'
+        ),
+    )
     args = parser.parse_args(argv)
 
-    protected = set(args.branch or ('master',))
-    return int(is_on_branch(protected))
+    protected = frozenset(args.branch or ('master',))
+    patterns = frozenset(args.pattern or ())
+    return int(is_on_branch(protected, patterns))
 
 
 if __name__ == '__main__':
diff --git a/tests/no_commit_to_branch_test.py b/tests/no_commit_to_branch_test.py
index e978ba2..a2ab1f1 100644
--- a/tests/no_commit_to_branch_test.py
+++ b/tests/no_commit_to_branch_test.py
@@ -44,6 +44,19 @@
         assert main(('--branch', 'b1', '--branch', 'b2'))
 
 
+def test_branch_pattern_fail(temp_git_dir):
+    with temp_git_dir.as_cwd():
+        cmd_output('git', 'checkout', '-b', 'another/branch')
+        assert is_on_branch(set(), {'another/.*'}) is True
+
+
+@pytest.mark.parametrize('branch_name', ('master', 'another/branch'))
+def test_branch_pattern_multiple_branches_fail(temp_git_dir, branch_name):
+    with temp_git_dir.as_cwd():
+        cmd_output('git', 'checkout', '-b', branch_name)
+        assert main(('--branch', 'master', '--pattern', 'another/.*'))
+
+
 def test_main_default_call(temp_git_dir):
     with temp_git_dir.as_cwd():
         cmd_output('git', 'checkout', '-b', 'anotherbranch')