Merge pull request #274 from pre-commit/check_yaml_unsafe

Add an `--unsafe` option to `check-yaml`

Author: asottile

README.md

@@ -52,6 +52,11 @@ Add this to your `.pre-commit-config.yaml`
 - `check-yaml` - Attempts to load all yaml files to verify syntax.
     - `--allow-multiple-documents` - allow yaml files which use the
       [multi-document syntax](http://www.yaml.org/spec/1.2/spec.html#YAML)
+    - `--unsafe` - Instaed of loading the files, simply parse them for syntax.
+      A syntax-only check enables extensions and unsafe constructs which would
+      otherwise be forbidden.  Using this option removes all guarantees of
+      portability to other yaml implementations.
+      Implies `--allow-multiple-documents`.
 - `debug-statements` - Check for pdb / ipdb / pudb statements in code.
 - `detect-aws-credentials` - Checks for the existence of AWS secrets that you
   have set up with the AWS CLI.

pre_commit_hooks/check_yaml.py

@@ -1,6 +1,7 @@
 from __future__ import print_function
 
 import argparse
+import collections
 import sys
 
 import yaml
@@ -11,24 +12,52 @@
     Loader = yaml.SafeLoader
 
 
+def _exhaust(gen):
+    for _ in gen:
+        pass
+
+
+def _parse_unsafe(*args, **kwargs):
+    _exhaust(yaml.parse(*args, **kwargs))
+
+
 def _load_all(*args, **kwargs):
-    # need to exhaust the generator
-    return tuple(yaml.load_all(*args, **kwargs))
+    _exhaust(yaml.load_all(*args, **kwargs))
+
+
+Key = collections.namedtuple('Key', ('multi', 'unsafe'))
+LOAD_FNS = {
+    Key(multi=False, unsafe=False): yaml.load,
+    Key(multi=False, unsafe=True): _parse_unsafe,
+    Key(multi=True, unsafe=False): _load_all,
+    Key(multi=True, unsafe=True): _parse_unsafe,
+}
 
 
 def check_yaml(argv=None):
     parser = argparse.ArgumentParser()
     parser.add_argument(
-        '-m', '--allow-multiple-documents', dest='yaml_load_fn',
-        action='store_const', const=_load_all, default=yaml.load,
+        '-m', '--multi', '--allow-multiple-documents', action='store_true',
+    )
+    parser.add_argument(
+        '--unsafe', action='store_true',
+        help=(
+            'Instead of loading the files, simply parse them for syntax.  '
+            'A syntax-only check enables extensions and unsafe contstructs '
+            'which would otherwise be forbidden.  Using this option removes '
+            'all guarantees of portability to other yaml implementations.  '
+            'Implies --allow-multiple-documents'
+        ),
     )
     parser.add_argument('filenames', nargs='*', help='Yaml filenames to check.')
     args = parser.parse_args(argv)
 
+    load_fn = LOAD_FNS[Key(multi=args.multi, unsafe=args.unsafe)]
+
     retval = 0
     for filename in args.filenames:
         try:
-            args.yaml_load_fn(open(filename), Loader=Loader)
+            load_fn(open(filename), Loader=Loader)
         except yaml.YAMLError as exc:
             print(exc)
             retval = 1

tests/check_yaml_test.py

@@ -22,7 +22,7 @@ def test_check_yaml_allow_multiple_documents(tmpdir):
     f = tmpdir.join('test.yaml')
     f.write('---\nfoo\n---\nbar\n')
 
-    # should failw without the setting
+    # should fail without the setting
     assert check_yaml((f.strpath,))
 
     # should pass when we allow multiple documents
@@ -33,3 +33,22 @@ def test_fails_even_with_allow_multiple_documents(tmpdir):
     f = tmpdir.join('test.yaml')
     f.write('[')
     assert check_yaml(('--allow-multiple-documents', f.strpath))
+
+
+def test_check_yaml_unsafe(tmpdir):
+    f = tmpdir.join('test.yaml')
+    f.write(
+        'some_foo: !vault |\n'
+        '    $ANSIBLE_VAULT;1.1;AES256\n'
+        '    deadbeefdeadbeefdeadbeef\n',
+    )
+    # should fail "safe" check
+    assert check_yaml((f.strpath,))
+    # should pass when we allow unsafe documents
+    assert not check_yaml(('--unsafe', f.strpath))
+
+
+def test_check_yaml_unsafe_still_fails_on_syntax_errors(tmpdir):
+    f = tmpdir.join('test.yaml')
+    f.write('[')
+    assert check_yaml(('--unsafe', f.strpath))