Merge pull request #244 from pre-commit/yaml_multi_document

Add an --allow-multiple-documents option to check-yaml
diff --git a/README.md b/README.md
index c4f9dc1..79d28fb 100644
--- a/README.md
+++ b/README.md
@@ -43,6 +43,8 @@
 - `check-vcs-permalinks` - Ensures that links to vcs websites are permalinks.
 - `check-xml` - Attempts to load all xml files to verify syntax.
 - `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)
 - `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.
diff --git a/pre_commit_hooks/check_yaml.py b/pre_commit_hooks/check_yaml.py
index cc9a614..e9bb8f0 100644
--- a/pre_commit_hooks/check_yaml.py
+++ b/pre_commit_hooks/check_yaml.py
@@ -11,15 +11,24 @@
     Loader = yaml.SafeLoader
 
 
+def _load_all(*args, **kwargs):
+    # need to exhaust the generator
+    return tuple(yaml.load_all(*args, **kwargs))
+
+
 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,
+    )
     parser.add_argument('filenames', nargs='*', help='Yaml filenames to check.')
     args = parser.parse_args(argv)
 
     retval = 0
     for filename in args.filenames:
         try:
-            yaml.load(open(filename), Loader=Loader)
+            args.yaml_load_fn(open(filename), Loader=Loader)
         except yaml.YAMLError as exc:
             print(exc)
             retval = 1
diff --git a/tests/check_yaml_test.py b/tests/check_yaml_test.py
index 73d6593..de3b383 100644
--- a/tests/check_yaml_test.py
+++ b/tests/check_yaml_test.py
@@ -1,3 +1,6 @@
+from __future__ import absolute_import
+from __future__ import unicode_literals
+
 import pytest
 
 from pre_commit_hooks.check_yaml import check_yaml
@@ -13,3 +16,20 @@
 def test_check_yaml(filename, expected_retval):
     ret = check_yaml([get_resource_path(filename)])
     assert ret == expected_retval
+
+
+def test_check_yaml_allow_multiple_documents(tmpdir):
+    f = tmpdir.join('test.yaml')
+    f.write('---\nfoo\n---\nbar\n')
+
+    # should failw without the setting
+    assert check_yaml((f.strpath,))
+
+    # should pass when we allow multiple documents
+    assert not check_yaml(('--allow-multiple-documents', f.strpath))
+
+
+def test_fails_even_with_allow_multiple_documents(tmpdir):
+    f = tmpdir.join('test.yaml')
+    f.write('[')
+    assert check_yaml(('--allow-multiple-documents', f.strpath))