From f72210df1f179b26d1781f6555aba8d731d11cd0 Mon Sep 17 00:00:00 2001
From: Martin Migasiewicz
 <71527391+martinmigasiewicz-tomtom@users.noreply.github.com>
Date: Tue, 21 May 2024 10:02:54 +0200
Subject: [PATCH] feat: support multi document yaml files

integrated from https://github.com/python-jsonschema/check-jsonschema/pull/223
---
 src/check_jsonschema/checker.py               | 10 ++++++----
 src/check_jsonschema/instance_loader.py       |  3 +++
 src/check_jsonschema/parsers/yaml.py          |  9 +++++++++
 src/check_jsonschema/schema_loader/readers.py |  5 ++++-
 4 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/src/check_jsonschema/checker.py b/src/check_jsonschema/checker.py
index e3e1b9f1e..9fefb819a 100644
--- a/src/check_jsonschema/checker.py
+++ b/src/check_jsonschema/checker.py
@@ -68,11 +68,13 @@ def _build_result(self) -> CheckResult:
             if isinstance(data, ParseError):
                 result.record_parse_error(path, data)
             else:
-                validator = self.get_validator(path, data)
                 passing = True
-                for err in validator.iter_errors(data):
-                    result.record_validation_error(path, err)
-                    passing = False
+                data_list = data if isinstance(data, list) else [data]
+                for data in data_list:
+                    validator = self.get_validator(path, data)
+                    for err in validator.iter_errors(data):
+                        result.record_validation_error(path, err)
+                        passing = False
                 if passing:
                     result.record_validation_success(path)
         return result
diff --git a/src/check_jsonschema/instance_loader.py b/src/check_jsonschema/instance_loader.py
index 2d0651c4f..547cc99a3 100644
--- a/src/check_jsonschema/instance_loader.py
+++ b/src/check_jsonschema/instance_loader.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import io
+import types
 import typing as t
 
 from check_jsonschema.cli.param_types import CustomLazyFile
@@ -51,6 +52,8 @@ def iter_files(self) -> t.Iterator[tuple[str, ParseError | t.Any]]:
                 except ParseError as err:
                     data = err
                 else:
+                    if isinstance(data, types.GeneratorType):
+                        data = list(data)
                     data = self._data_transform(data)
             finally:
                 file.close()
diff --git a/src/check_jsonschema/parsers/yaml.py b/src/check_jsonschema/parsers/yaml.py
index d2780d77c..61befea15 100644
--- a/src/check_jsonschema/parsers/yaml.py
+++ b/src/check_jsonschema/parsers/yaml.py
@@ -64,6 +64,15 @@ def load(stream: t.IO[bytes]) -> t.Any:
                 try:
                     data = impl.load(stream_bytes)
                 except ruamel.yaml.YAMLError as e:
+                    if isinstance(
+                        e, ruamel.yaml.composer.ComposerError
+                    ) and "expected a single document in the stream" in str(e):
+                        try:
+                            data = impl.load_all(stream_bytes)
+                        except ruamel.yaml.YAMLError as e:
+                            lasterr = e
+                        else:
+                            break
                     lasterr = e
                 else:
                     break
diff --git a/src/check_jsonschema/schema_loader/readers.py b/src/check_jsonschema/schema_loader/readers.py
index 907ce6936..d58d8d71e 100644
--- a/src/check_jsonschema/schema_loader/readers.py
+++ b/src/check_jsonschema/schema_loader/readers.py
@@ -3,6 +3,7 @@
 import io
 import json
 import sys
+import types
 import typing as t
 
 import ruamel.yaml
@@ -28,7 +29,9 @@ def _run_load_callback(schema_location: str, callback: t.Callable) -> dict:
     # only local loads can raise the YAMLError, but catch for both cases for simplicity
     except (ValueError, ruamel.yaml.error.YAMLError) as e:
         raise SchemaParseError(schema_location) from e
-    if not isinstance(schema, dict):
+    if isinstance(schema, types.GeneratorType):
+        schema = list(schema)
+    if not isinstance(schema, dict) and not isinstance(schema, list):
         raise SchemaParseError(schema_location)
     return schema