Fix: #182. Improve yaml support.

datamodel_code_generator/__init__.py

@@ -74,14 +74,6 @@ def inner(cls: Type[T]) -> Type[T]:
     return inner
 
 
-def load_json_or_yaml(text: str) -> Dict[Any, Any]:
-    try:
-        data = json.loads(text)
-    except JSONDecodeError:
-        data = yaml.safe_load(text)
-    return data
-
-
 @contextlib.contextmanager
 def chdir(path: Optional[Path]) -> Iterator[None]:
     """Changes working directory and returns to previous on exit."""
@@ -98,7 +90,7 @@ def chdir(path: Optional[Path]) -> Iterator[None]:
 
 
 def is_openapi(text: str) -> bool:
-    return 'openapi' in load_json_or_yaml(text)
+    return 'openapi' in yaml.safe_load(text)
 
 
 class InputFileType(Enum):

datamodel_code_generator/parser/jsonschema.py

@@ -14,6 +14,7 @@
     Union,
 )
 
+import yaml
 from pydantic import BaseModel, Field, validator
 
 from datamodel_code_generator import snooper_to_methods
@@ -560,8 +561,6 @@ def parse_ref(self, obj: JsonSchemaObject, path: List[str]) -> None:
                     # Local Reference – $ref: '#/definitions/myElement'
                     pass
                 else:
-                    import yaml
-
                     relative_path, object_path = obj.ref.split('#/')
                     remote_object: Optional[
                         Dict[str, Any]
@@ -633,7 +632,7 @@ def parse_raw_obj(self, name: str, raw: Dict[str, Any], path: List[str]) -> None
         self.parse_ref(obj, path)
 
     def parse_raw(self) -> None:
-        raw_obj: Dict[Any, Any] = json.loads(self.text)  # type: ignore
+        raw_obj: Dict[Any, Any] = yaml.safe_load(self.text)  # type: ignore
         obj_name = raw_obj.get('title', 'Model')
         obj_name = self.model_resolver.add([], obj_name, unique=False).name
         self.parse_raw_obj(obj_name, raw_obj, [])

datamodel_code_generator/parser/openapi.py

@@ -1,6 +1,7 @@
 from typing import Any, Dict
+import yaml
 
-from datamodel_code_generator import load_json_or_yaml, snooper_to_methods
+from datamodel_code_generator import snooper_to_methods
 from datamodel_code_generator.parser.jsonschema import JsonSchemaParser
 
 
@@ -15,7 +16,7 @@ def parse_raw(self) -> None:
             )
             components: Dict[str, Any] = base_parser.specification['components']
         else:
-            components = load_json_or_yaml(self.text)['components']  # type: ignore
+            components = yaml.safe_load(self.text)['components']  # type: ignore
 
         for obj_name, raw_obj in components[
             'schemas'

tests/data/expected/parser/openapi/openapi_parser_parse_remote_ref/with_import_format.py

@@ -0,0 +1,13 @@
+from __future__ import annotations
+
+from typing import Optional
+
+from pydantic import AnyUrl, BaseModel, conint
+
+
+class Problem(BaseModel):
+    detail: Optional[str] = None
+    instance: Optional[AnyUrl] = None
+    status: Optional[conint(ge=100, le=600, lt=1)] = None
+    title: Optional[str] = None
+    type: Optional[AnyUrl] = 'about:blank'

tests/data/jsonschema/nested_array.json.snapshot

@@ -0,0 +1,8 @@
+class BoundingBox(BaseModel):
+    type: str
+    coordinates: List[Union[float, str]]
+
+
+class Model(BaseModel):
+    bounding_box: Optional[BoundingBox] = None
+    attributes: Optional[Dict[str, Any]] = None

tests/data/jsonschema/oneof.json

@@ -0,0 +1,20 @@
+{
+    "properties": {
+        "item": {
+            "properties": {
+                "timeout": {
+                    "oneOf": [
+                        {
+                            "type": "string"
+                        },
+                        {
+                            "type": "integer"
+                        }
+                    ]
+                }
+            },
+            "type": "object"
+        }
+    }
+}
+

tests/data/jsonschema/oneof.json.snapshot

@@ -0,0 +1,6 @@
+class Item(BaseModel):
+    timeout: Optional[Union[str, int]] = None
+
+
+class OnOfObject(BaseModel):
+    item: Optional[Item] = None

tests/data/openapi/definitions.yaml

@@ -0,0 +1,42 @@
+schemas:
+  Problem:
+    properties:
+      detail:
+        description: |
+          A human readable explanation specific to this occurrence of the
+          problem. You MUST NOT expose internal informations, personal
+          data or implementation details through this field.
+        example: Request took too long to complete.
+        type: string
+      instance:
+        description: |
+          An absolute URI that identifies the specific occurrence of the problem.
+          It may or may not yield further information if dereferenced.
+        format: uri
+        type: string
+      status:
+        description: |
+          The HTTP status code generated by the origin server for this occurrence
+          of the problem.
+        example: 503
+        exclusiveMaximum: true
+        format: int32
+        maximum: 600
+        minimum: 100
+        type: integer
+      title:
+        description: |
+          A short, summary of the problem type. Written in english and readable
+          for engineers (usually not suited for non technical stakeholders and
+          not localized); example: Service Unavailable
+        type: string
+      type:
+        default: about:blank
+        description: |
+          An absolute URI that identifies the problem type.  When dereferenced,
+          it SHOULD provide human-readable documentation for the problem type
+          (e.g., using HTML).
+        example: https://tools.ietf.org/html/rfc7231#section-6.6.4
+        format: uri
+        type: string
+    type: object

tests/data/openapi/refs.yaml

@@ -0,0 +1,5 @@
+openapi: 3.0.1
+components:
+  schemas:
+    Problem:
+      $ref: "https://teamdigitale.github.io/openapi/0.0.6/definitions.yaml#/schemas/Problem"

tests/parser/test_jsonschema.py

@@ -16,6 +16,7 @@
 )
 
 DATA_PATH: Path = Path(__file__).parents[1] / 'data' / 'jsonschema'
+OAS_PATH: Path = Path(__file__).parents[1] / 'data' / 'yaml'
 
 
 @pytest.mark.parametrize(
@@ -260,24 +261,8 @@ def test_parse_any_root_object(source_obj, generated_classes):
     'source_obj,generated_classes',
     [
         (
-            {
-                "properties": {
-                    "item": {
-                        "properties": {
-                            "timeout": {
-                                "oneOf": [{"type": "string"}, {"type": "integer"}]
-                            }
-                        },
-                        "type": "object",
-                    }
-                }
-            },
-            """class Item(BaseModel):
-    timeout: Optional[Union[str, int]] = None
-
-
-class OnOfObject(BaseModel):
-    item: Optional[Item] = None""",
+            yaml.safe_load((DATA_PATH / 'oneof.json').read_text()),
+            (DATA_PATH / 'oneof.json.snapshot').read_text(),
         )
     ],
 )
@@ -345,13 +330,7 @@ def test_parse_nested_array():
     parser.parse()
     assert (
         dump_templates(list(parser.results))
-        == """\
-class BoundingBox(BaseModel):
-    type: str
-    coordinates: List[Union[float, str]]
+        == (DATA_PATH / 'nested_array.json.snapshot').read_text()
+    )
 
 
-class Model(BaseModel):
-    bounding_box: Optional[BoundingBox] = None
-    attributes: Optional[Dict[str, Any]] = None"""
-    )

tests/parser/test_openapi.py

@@ -527,3 +527,22 @@ def test_openapi_parser_parse_array_enum(with_import, format_, base_class):
         parser.parse(with_import=with_import, format_=format_)
         == expected_file.read_text()
     )
+
+@pytest.mark.parametrize(
+    'with_import, format_, base_class', [(True, True, None,),],
+)
+def test_openapi_parser_parse_remote_ref(with_import, format_, base_class):
+    parser = OpenAPIParser(
+        BaseModel,
+        CustomRootType,
+        base_class=base_class,
+        text=(DATA_PATH / 'refs.yaml').read_text(),
+    )
+    expected_file = get_expected_file(
+        'openapi_parser_parse_remote_ref', with_import, format_, base_class
+    )
+
+    assert (
+        parser.parse(with_import=with_import, format_=format_)
+        == expected_file.read_text()
+    )