loader: support json and yaml separately

json is not exactly a subset of yaml, as yaml do not support tabs
there are maybe other subtleties, so we just use content-type to figure out how to parse

Signed-off-by: Pierre Tardy <[email protected]>

Author: tardyp

openapi_python_client/__init__.py

@@ -1,5 +1,7 @@
 """ Generate modern Python clients from OpenAPI """
 
+import json
+import mimetypes
 import shutil
 import subprocess
 import sys
@@ -361,21 +363,40 @@ def update_existing_client(
     return project.update()
 
 
+def _load_yaml_or_json(data: bytes, content_type: Optional[str]) -> Union[Dict[str, Any], GeneratorError]:
+    if content_type == "application/json":
+        try:
+            return json.loads(data.decode())
+        except ValueError as err:
+            return GeneratorError(header="Invalid JSON from provided source: {}".format(str(err)))
+    else:
+        try:
+            return yaml.safe_load(data)
+        except yaml.YAMLError as err:
+            return GeneratorError(header="Invalid YAML from provided source: {}".format(str(err)))
+
+
 def _get_document(*, url: Optional[str], path: Optional[Path]) -> Union[Dict[str, Any], GeneratorError]:
     yaml_bytes: bytes
+    content_type: Optional[str]
     if url is not None and path is not None:
         return GeneratorError(header="Provide URL or Path, not both.")
     if url is not None:
         try:
             response = httpx.get(url)
             yaml_bytes = response.content
+            if "content-type" in response.headers:
+                content_type = response.headers["content-type"].split(";")[0]
+            else:
+                content_type = mimetypes.guess_type(url, strict=True)[0]
+
         except (httpx.HTTPError, httpcore.NetworkError):
             return GeneratorError(header="Could not get OpenAPI document from provided URL")
     elif path is not None:
         yaml_bytes = path.read_bytes()
+        content_type = mimetypes.guess_type(path.as_uri(), strict=True)[0]
+
     else:
         return GeneratorError(header="No URL or Path provided")
-    try:
-        return yaml.safe_load(yaml_bytes)
-    except yaml.YAMLError:
-        return GeneratorError(header="Invalid YAML from provided source")
+
+    return _load_yaml_or_json(yaml_bytes, content_type)

tests/test___init__.py

@@ -218,7 +218,7 @@ def test__get_document_path_no_url(self, mocker):
 
     def test__get_document_bad_yaml(self, mocker):
         get = mocker.patch("httpx.get")
-        loads = mocker.patch("yaml.safe_load", side_effect=yaml.YAMLError)
+        loads = mocker.patch("yaml.safe_load", side_effect=yaml.YAMLError("error line 2"))
 
         from openapi_python_client import _get_document
 
@@ -228,7 +228,44 @@ def test__get_document_bad_yaml(self, mocker):
         get.assert_not_called()
         path.read_bytes.assert_called_once()
         loads.assert_called_once_with(path.read_bytes())
-        assert result == GeneratorError(header="Invalid YAML from provided source")
+        assert result == GeneratorError(header="Invalid YAML from provided source: error line 2")
+
+    def test__get_document_json(self, mocker):
+        class FakeResponse:
+            content = b'{\n\t"foo": "bar"}'
+            headers = {"content-type": "application/json; encoding=utf8"}
+
+        get = mocker.patch("httpx.get", return_value=FakeResponse())
+        yaml_loads = mocker.patch("yaml.safe_load")
+        json_result = mocker.MagicMock()
+        json_loads = mocker.patch("json.loads", return_value=json_result)
+
+        from openapi_python_client import _get_document
+
+        url = mocker.MagicMock()
+        result = _get_document(url=url, path=None)
+
+        get.assert_called_once()
+        json_loads.assert_called_once_with(FakeResponse.content.decode())
+        assert result == json_result
+
+    def test__get_document_bad_json(self, mocker):
+        class FakeResponse:
+            content = b'{"foo"}'
+            headers = {"content-type": "application/json; encoding=utf8"}
+
+        get = mocker.patch("httpx.get", return_value=FakeResponse())
+        json_result = mocker.MagicMock()
+
+        from openapi_python_client import _get_document
+
+        url = mocker.MagicMock()
+        result = _get_document(url=url, path=None)
+
+        get.assert_called_once()
+        assert result == GeneratorError(
+            header="Invalid JSON from provided source: " "Expecting ':' delimiter: line 1 column 7 (char 6)"
+        )
 
 
 def make_project(**kwargs):