Merge pull request #1395 from pre-commit/argument_length

asottile · web-flow · commit 171fd18ba5cf · 2020-04-15T13:46:13.000-07:00
validate argument length as part of hook-impl
diff --git a/pre_commit/commands/hook_impl.py b/pre_commit/commands/hook_impl.py
@@ -147,15 +147,44 @@ def _pre_push_ns(
     return None
 
 
+_EXPECTED_ARG_LENGTH_BY_HOOK = {
+    'commit-msg': 1,
+    'post-checkout': 3,
+    'pre-commit': 0,
+    'pre-merge-commit': 0,
+    'pre-push': 2,
+}
+
+
+def _check_args_length(hook_type: str, args: Sequence[str]) -> None:
+    if hook_type == 'prepare-commit-msg':
+        if len(args) < 1 or len(args) > 3:
+            raise SystemExit(
+                f'hook-impl for {hook_type} expected 1, 2, or 3 arguments '
+                f'but got {len(args)}: {args}',
+            )
+    elif hook_type in _EXPECTED_ARG_LENGTH_BY_HOOK:
+        expected = _EXPECTED_ARG_LENGTH_BY_HOOK[hook_type]
+        if len(args) != expected:
+            arguments_s = 'argument' if expected == 1 else 'arguments'
+            raise SystemExit(
+                f'hook-impl for {hook_type} expected {expected} {arguments_s} '
+                f'but got {len(args)}: {args}',
+            )
+    else:
+        raise AssertionError(f'unexpected hook type: {hook_type}')
+
+
 def _run_ns(
         hook_type: str,
         color: bool,
         args: Sequence[str],
         stdin: bytes,
 ) -> Optional[argparse.Namespace]:
+    _check_args_length(hook_type, args)
     if hook_type == 'pre-push':
         return _pre_push_ns(color, args, stdin)
-    elif hook_type in {'prepare-commit-msg', 'commit-msg'}:
+    elif hook_type in {'commit-msg', 'prepare-commit-msg'}:
         return _ns(hook_type, color, commit_msg_filename=args[0])
     elif hook_type in {'pre-merge-commit', 'pre-commit'}:
         return _ns(hook_type, color)
diff --git a/tests/commands/hook_impl_test.py b/tests/commands/hook_impl_test.py
@@ -89,6 +89,51 @@ def call(*_, **__):
             call()
 
 
+@pytest.mark.parametrize(
+    ('hook_type', 'args'),
+    (
+        ('pre-commit', []),
+        ('pre-merge-commit', []),
+        ('pre-push', ['branch_name', 'remote_name']),
+        ('commit-msg', ['.git/COMMIT_EDITMSG']),
+        ('post-checkout', ['old_head', 'new_head', '1']),
+        # multiple choices for commit-editmsg
+        ('prepare-commit-msg', ['.git/COMMIT_EDITMSG']),
+        ('prepare-commit-msg', ['.git/COMMIT_EDITMSG', 'message']),
+        ('prepare-commit-msg', ['.git/COMMIT_EDITMSG', 'commit', 'deadbeef']),
+    ),
+)
+def test_check_args_length_ok(hook_type, args):
+    hook_impl._check_args_length(hook_type, args)
+
+
+def test_check_args_length_error_too_many_plural():
+    with pytest.raises(SystemExit) as excinfo:
+        hook_impl._check_args_length('pre-commit', ['run', '--all-files'])
+    msg, = excinfo.value.args
+    assert msg == (
+        'hook-impl for pre-commit expected 0 arguments but got 2: '
+        "['run', '--all-files']"
+    )
+
+
+def test_check_args_length_error_too_many_singluar():
+    with pytest.raises(SystemExit) as excinfo:
+        hook_impl._check_args_length('commit-msg', [])
+    msg, = excinfo.value.args
+    assert msg == 'hook-impl for commit-msg expected 1 argument but got 0: []'
+
+
+def test_check_args_length_prepare_commit_msg_error():
+    with pytest.raises(SystemExit) as excinfo:
+        hook_impl._check_args_length('prepare-commit-msg', [])
+    msg, = excinfo.value.args
+    assert msg == (
+        'hook-impl for prepare-commit-msg expected 1, 2, or 3 arguments '
+        'but got 0: []'
+    )
+
+
 def test_run_ns_pre_commit():
     ns = hook_impl._run_ns('pre-commit', True, (), b'')
     assert ns is not None
