From b8a27e7d762bc3698a850eab7f8c8bff96eac9b9 Mon Sep 17 00:00:00 2001
From: XieJiSS <c141028@gmail.com>
Date: Mon, 16 Dec 2024 00:15:14 +0800
Subject: [PATCH] feat: support InputFileType.Json and InputFileType.Dict
 (#2165)

* feat: support InputFileType.Json and InputFileType.Dict

* fix: support Mapping as input_ of generate()

* chore: use read_text for InputFileType.Json for consistency

* fix: fixed failing test cases

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix: incorrect handling of direct dict input

* test: add test for direct dict input

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix: remove dead code

* Add else pattern

---------

Co-authored-by: Koudai Aono <koxudaxi@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: XieJiSS <XieJiSS@users.noreply.github.com>
---
 datamodel_code_generator/__init__.py          | 29 +++++++++++++++++--
 tests/data/expected/main/direct_input_dict.py | 16 ++++++++++
 tests/main/test_main_general.py               | 21 ++++++++++++++
 3 files changed, 63 insertions(+), 3 deletions(-)
 create mode 100644 tests/data/expected/main/direct_input_dict.py

diff --git a/datamodel_code_generator/__init__.py b/datamodel_code_generator/__init__.py
index 6ebfab2e..7445e626 100644
--- a/datamodel_code_generator/__init__.py
+++ b/datamodel_code_generator/__init__.py
@@ -231,7 +231,7 @@ def get_first_file(path: Path) -> Path:  # pragma: no cover
 
 
 def generate(
-    input_: Union[Path, str, ParseResult],
+    input_: Union[Path, str, ParseResult, Mapping[str, Any]],
     *,
     input_filename: Optional[str] = None,
     input_file_type: InputFileType = InputFileType.Auto,
@@ -354,6 +354,8 @@ def generate(
         parser_class = JsonSchemaParser
 
         if input_file_type in RAW_DATA_TYPES:
+            import json
+
             try:
                 if isinstance(input_, Path) and input_.is_dir():  # pragma: no cover
                     raise Error(f'Input must be a file for {input_file_type}')
@@ -372,15 +374,33 @@ def get_header_and_first_line(csv_file: IO[str]) -> Dict[str, Any]:
                         import io
 
                         obj = get_header_and_first_line(io.StringIO(input_text))
-                else:
+                elif input_file_type == InputFileType.Yaml:
                     obj = load_yaml(
                         input_.read_text(encoding=encoding)  # type: ignore
                         if isinstance(input_, Path)
                         else input_text
                     )
+                elif input_file_type == InputFileType.Json:
+                    obj = json.loads(
+                        input_.read_text(encoding=encoding)  # type: ignore
+                        if isinstance(input_, Path)
+                        else input_text
+                    )
+                elif input_file_type == InputFileType.Dict:
+                    import ast
+
+                    # Input can be a dict object stored in a python file
+                    obj = (
+                        ast.literal_eval(
+                            input_.read_text(encoding=encoding)  # type: ignore
+                        )
+                        if isinstance(input_, Path)
+                        else input_
+                    )
+                else:  # pragma: no cover
+                    raise Error(f'Unsupported input file type: {input_file_type}')
             except:  # noqa
                 raise Error('Invalid file format')
-            import json
 
             from genson import SchemaBuilder
 
@@ -490,6 +510,9 @@ def get_header_and_first_line(csv_file: IO[str]) -> Dict[str, Any]:
             input_filename = '<stdin>'
         elif isinstance(input_, ParseResult):
             input_filename = input_.geturl()
+        elif input_file_type == InputFileType.Dict:
+            # input_ might be a dict object provided directly, and missing a name field
+            input_filename = getattr(input_, 'name', '<dict>')
         else:
             input_filename = input_.name
     if not results:
diff --git a/tests/data/expected/main/direct_input_dict.py b/tests/data/expected/main/direct_input_dict.py
new file mode 100644
index 00000000..532496dd
--- /dev/null
+++ b/tests/data/expected/main/direct_input_dict.py
@@ -0,0 +1,16 @@
+# generated by datamodel-codegen:
+#   filename:  <dict>
+#   timestamp: 2024-12-14T00:00:00+00:00
+
+from __future__ import annotations
+
+from pydantic import BaseModel
+
+
+class Bar(BaseModel):
+    baz: int
+
+
+class Model(BaseModel):
+    foo: int
+    bar: Bar
diff --git a/tests/main/test_main_general.py b/tests/main/test_main_general.py
index 60fd4855..bfb42fee 100644
--- a/tests/main/test_main_general.py
+++ b/tests/main/test_main_general.py
@@ -6,6 +6,9 @@
 from freezegun import freeze_time
 
 from datamodel_code_generator import (
+    DataModelType,
+    InputFileType,
+    generate,
     snooper_to_methods,
 )
 from datamodel_code_generator.__main__ import Exit, main
@@ -86,3 +89,21 @@ def test_space_and_special_characters_dict():
             output_file.read_text()
             == (EXPECTED_MAIN_PATH / 'space_and_special_characters_dict.py').read_text()
         )
+
+
+@freeze_time('2024-12-14')
+def test_direct_input_dict():
+    with TemporaryDirectory() as output_dir:
+        output_file = Path(output_dir) / 'output.py'
+        generate(
+            {'foo': 1, 'bar': {'baz': 2}},
+            input_file_type=InputFileType.Dict,
+            input_filename=None,
+            output=output_file,
+            output_model_type=DataModelType.PydanticV2BaseModel,
+            snake_case_field=True,
+        )
+        assert (
+            output_file.read_text()
+            == (EXPECTED_MAIN_PATH / 'direct_input_dict.py').read_text()
+        )