Repository: koxudaxi/datamodel-code-generator

Issue:
Ability to add custom imports
**Is your feature request related to a problem? Please describe.**
Can't customize my template to add new imports

**Describe the solution you'd like**
New argument `additional_imports` in `generate` function which adds additional imports to final rendered template

**Describe alternatives you've considered**
-

**Additional context**
-

Discussion:
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/koxudaxi/datamodel-code-generator/pull/1422?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Koudai+Aono) Report
Patch coverage: **`99.48`**% and project coverage change: **`-0.46`** :warning:
> Comparison is base [(`e9b6edf`)](https://app.codecov.io/gh/koxudaxi/datamodel-code-generator/commit/e9b6edfc2e2ace195cb465e2e980fb78c22c3fd9?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Koudai+Aono) 100.00% compared to head [(`8a106b4`)](https://app.codecov.io/gh/koxudaxi/datamodel-code-generator/pull/1422?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Koudai+Aono) 99.54%.

> :exclamation: Current head 8a106b4 differs from pull request most recent head 6225b8f. Consider uploading reports for the commit 6225b8f to get more accurate results

<details><summary>Additional details and impacted files</summary>


```diff
@@             Coverage Diff             @@
##            master    #1422      +/-   ##
===========================================
- Coverage   100.00%   99.54%   -0.46%     
===========================================
  Files           11       32      +21     
  Lines         1020     3547    +2527     
  Branches       201      833     +632     
===========================================
+ Hits          1020     3531    +2511     
- Misses           0       16      +16     
```

| Flag | Coverage Δ | |
|---|---|---|
| unittests | `99.18% <99.06%> (?)` | |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Koudai+Aono#carryforward-flags-in-the-pull-request-comment) to find out more.

| [Impacted Files](https://app.codecov.io/gh/koxudaxi/datamodel-code-generator/pull/1422?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Koudai+Aono) | Coverage Δ | |
|---|---|---|
| [datamodel\_code\_generator/util.py](https://app.codecov.io/gh/koxudaxi/datamodel-code-generator/pull/1422?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Koudai+Aono#diff-ZGF0YW1vZGVsX2NvZGVfZ2VuZXJhdG9yL3V0aWwucHk=) | `75.00% <75.00%> (ø)` | |
| [datamodel\_code\_generator/\_\_init\_\_.py](https://app.codecov.io/gh/koxudaxi/datamodel-code-generator/pull/1422?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Koudai+Aono#diff-ZGF0YW1vZGVsX2NvZGVfZ2VuZXJhdG9yL19faW5pdF9fLnB5) | `98.36% <97.34%> (ø)` | |
| [datamodel\_code\_generator/\_\_main\_\_.py](https://app.codecov.io/gh/koxudaxi/datamodel-code-generator/pull/1422?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Koudai+Aono#diff-ZGF0YW1vZGVsX2NvZGVfZ2VuZXJhdG9yL19fbWFpbl9fLnB5) | `100.00% <100.00%> (ø)` | |
| [datamodel\_code\_generator/format.py](https://app.codecov.io/gh/koxudaxi/datamodel-code-generator/pull/1422?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Koudai+Aono#diff-ZGF0YW1vZGVsX2NvZGVfZ2VuZXJhdG9yL2Zvcm1hdC5weQ==) | `100.00% <100.00%> (ø)` | |
| [datamodel\_code\_generator/http.py](https://app.codecov.io/gh/koxudaxi/datamodel-code-generator/pull/1422?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Koudai+Aono#diff-ZGF0YW1vZGVsX2NvZGVfZ2VuZXJhdG9yL2h0dHAucHk=) | `100.00% <100.00%> (ø)` | |
| [datamodel\_code\_generator/imports.py](https://app.codecov.io/gh/koxudaxi/datamodel-code-generator/pull/1422?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Koudai+Aono#diff-ZGF0YW1vZGVsX2NvZGVfZ2VuZXJhdG9yL2ltcG9ydHMucHk=) | `100.00% <100.00%> (ø)` | |
| [datamodel\_code\_generator/model/\_\_init\_\_.py](https://app.codecov.io/gh/koxudaxi/datamodel-code-generator/pull/1422?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Koudai+Aono#diff-ZGF0YW1vZGVsX2NvZGVfZ2VuZXJhdG9yL21vZGVsL19faW5pdF9fLnB5) | `100.00% <100.00%> (ø)` | |
| [datamodel\_code\_generator/model/base.py](https://app.codecov.io/gh/koxudaxi/datamodel-code-generator/pull/1422?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Koudai+Aono#diff-ZGF0YW1vZGVsX2NvZGVfZ2VuZXJhdG9yL21vZGVsL2Jhc2UucHk=) | `100.00% <100.00%> (ø)` | |
| [datamodel\_code\_generator/model/dataclass.py](https://app.codecov.io/gh/koxudaxi/datamodel-code-generator/pull/1422?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Koudai+Aono#diff-ZGF0YW1vZGVsX2NvZGVfZ2VuZXJhdG9yL21vZGVsL2RhdGFjbGFzcy5weQ==) | `100.00% <100.00%> (ø)` | |
| [datamodel\_code\_generator/model/enum.py](https://app.codecov.io/gh/koxudaxi/datamodel-code-generator/pull/1422?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Koudai+Aono#diff-ZGF0YW1vZGVsX2NvZGVfZ2VuZXJhdG9yL21vZGVsL2VudW0ucHk=) | `100.00% <100.00%> (ø)` | |
| ... and [22 more](https://app.codecov.io/gh/koxudaxi/datamodel-code-generator/pull/1422?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Koudai+Aono) | |


</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/koxudaxi/datamodel-code-generator/pull/1422?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Koudai+Aono).   
:loudspeaker: Do you have feedback about the report comment? [Let us know in this issue](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Koudai+Aono).


Patch:
diff --git a/datamodel_code_generator/__init__.py b/datamodel_code_generator/__init__.py
index a71ff67a..cceaa958 100644
--- a/datamodel_code_generator/__init__.py
+++ b/datamodel_code_generator/__init__.py
@@ -225,6 +225,7 @@ def generate(
     output_model_type: DataModelType = DataModelType.PydanticBaseModel,
     target_python_version: PythonVersion = PythonVersion.PY_37,
     base_class: str = DEFAULT_BASE_CLASS,
+    additional_imports: Optional[List[str]] = None,
     custom_template_dir: Optional[Path] = None,
     extra_template_data: Optional[DefaultDict[str, Dict[str, Any]]] = None,
     validation: bool = False,
@@ -369,6 +370,7 @@ def generate(
         data_model_field_type=data_model_types.field_model,
         data_type_manager_type=data_model_types.data_type_manager,
         base_class=base_class,
+        additional_imports=additional_imports,
         custom_template_dir=custom_template_dir,
         extra_template_data=extra_template_data,
         target_python_version=target_python_version,
diff --git a/datamodel_code_generator/parser/base.py b/datamodel_code_generator/parser/base.py
index 24bb31bc..e5a03bbe 100644
--- a/datamodel_code_generator/parser/base.py
+++ b/datamodel_code_generator/parser/base.py
@@ -314,6 +314,7 @@ class Parser(ABC):
         data_type_manager_type: Type[DataTypeManager] = pydantic_model.DataTypeManager,
         data_model_field_type: Type[DataModelFieldBase] = pydantic_model.DataModelField,
         base_class: Optional[str] = None,
+        additional_imports: Optional[List[str]] = None,
         custom_template_dir: Optional[Path] = None,
         extra_template_data: Optional[DefaultDict[str, Dict[str, Any]]] = None,
         target_python_version: PythonVersion = PythonVersion.PY_37,
@@ -380,7 +381,10 @@ class Parser(ABC):
         self.data_model_type: Type[DataModel] = data_model_type
         self.data_model_root_type: Type[DataModel] = data_model_root_type
         self.data_model_field_type: Type[DataModelFieldBase] = data_model_field_type
+
         self.imports: Imports = Imports()
+        self._append_additional_imports(additional_imports=additional_imports)
+
         self.base_class: Optional[str] = base_class
         self.target_python_version: PythonVersion = target_python_version
         self.results: List[DataModel] = []
@@ -504,6 +508,16 @@ class Parser(ABC):
                 ),
             )
 
+    def _append_additional_imports(
+        self, additional_imports: Optional[List[str]]
+    ) -> None:
+        if additional_imports is None:
+            additional_imports = []
+
+        for additional_import_string in additional_imports:
+            new_import = Import.from_full_path(additional_import_string)
+            self.imports.append(new_import)
+
     def _get_text_from_url(self, url: str) -> str:
         from datamodel_code_generator.http import get_body
 
diff --git a/datamodel_code_generator/parser/jsonschema.py b/datamodel_code_generator/parser/jsonschema.py
index 7a39b048..4812199a 100644
--- a/datamodel_code_generator/parser/jsonschema.py
+++ b/datamodel_code_generator/parser/jsonschema.py
@@ -371,6 +371,7 @@ class JsonSchemaParser(Parser):
         data_type_manager_type: Type[DataTypeManager] = pydantic_model.DataTypeManager,
         data_model_field_type: Type[DataModelFieldBase] = pydantic_model.DataModelField,
         base_class: Optional[str] = None,
+        additional_imports: Optional[List[str]] = None,
         custom_template_dir: Optional[Path] = None,
         extra_template_data: Optional[DefaultDict[str, Dict[str, Any]]] = None,
         target_python_version: PythonVersion = PythonVersion.PY_37,
@@ -432,6 +433,7 @@ class JsonSchemaParser(Parser):
             data_type_manager_type=data_type_manager_type,
             data_model_field_type=data_model_field_type,
             base_class=base_class,
+            additional_imports=additional_imports,
             custom_template_dir=custom_template_dir,
             extra_template_data=extra_template_data,
             target_python_version=target_python_version,
diff --git a/datamodel_code_generator/parser/openapi.py b/datamodel_code_generator/parser/openapi.py
index fc825d4e..c8032187 100644
--- a/datamodel_code_generator/parser/openapi.py
+++ b/datamodel_code_generator/parser/openapi.py
@@ -161,6 +161,7 @@ class OpenAPIParser(JsonSchemaParser):
         data_type_manager_type: Type[DataTypeManager] = pydantic_model.DataTypeManager,
         data_model_field_type: Type[DataModelFieldBase] = pydantic_model.DataModelField,
         base_class: Optional[str] = None,
+        additional_imports: Optional[List[str]] = None,
         custom_template_dir: Optional[Path] = None,
         extra_template_data: Optional[DefaultDict[str, Dict[str, Any]]] = None,
         target_python_version: PythonVersion = PythonVersion.PY_37,
@@ -223,6 +224,7 @@ class OpenAPIParser(JsonSchemaParser):
             data_type_manager_type=data_type_manager_type,
             data_model_field_type=data_model_field_type,
             base_class=base_class,
+            additional_imports=additional_imports,
             custom_template_dir=custom_template_dir,
             extra_template_data=extra_template_data,
             target_python_version=target_python_version,


Test patch:
diff --git a/tests/parser/test_base.py b/tests/parser/test_base.py
index a2e754aa..1a980891 100644
--- a/tests/parser/test_base.py
+++ b/tests/parser/test_base.py
@@ -214,3 +214,21 @@ class D(DataModel):
 
     def render(self) -> str:
         return self._data
+
+
+def test_additional_imports():
+    """Test that additional imports are inside imports container."""
+    new_parser = C(
+        source='',
+        additional_imports=['collections.deque'],
+    )
+    assert len(new_parser.imports) == 1
+    assert new_parser.imports['collections'] == {'deque'}
+
+
+def test_no_additional_imports():
+    """Test that not additional imports are not affecting imports container."""
+    new_parser = C(
+        source='',
+    )
+    assert len(new_parser.imports) == 0
diff --git a/tests/parser/test_jsonschema.py b/tests/parser/test_jsonschema.py
index 9088f6c3..bf50e2e1 100644
--- a/tests/parser/test_jsonschema.py
+++ b/tests/parser/test_jsonschema.py
@@ -439,3 +439,18 @@ def test_get_data_type_array(schema_types, result_types):
         ],
         is_optional='null' in schema_types,
     )
+
+
+def test_additional_imports():
+    """Test that additional imports are inside imports container."""
+    new_parser = JsonSchemaParser(source='', additional_imports=['collections.deque'])
+    assert len(new_parser.imports) == 1
+    assert new_parser.imports['collections'] == {'deque'}
+
+
+def test_no_additional_imports():
+    """Test that not additional imports are not affecting imports container."""
+    new_parser = JsonSchemaParser(
+        source='',
+    )
+    assert len(new_parser.imports) == 0
diff --git a/tests/parser/test_openapi.py b/tests/parser/test_openapi.py
index 88cd1512..c3647778 100644
--- a/tests/parser/test_openapi.py
+++ b/tests/parser/test_openapi.py
@@ -729,3 +729,18 @@ def test_openapi_parser_with_query_parameters():
             / 'output.py'
         ).read_text()
     )
+
+
+def test_additional_imports():
+    """Test that additional imports are inside imports container."""
+    new_parser = OpenAPIParser(source='', additional_imports=['collections.deque'])
+    assert len(new_parser.imports) == 1
+    assert new_parser.imports['collections'] == {'deque'}
+
+
+def test_no_additional_imports():
+    """Test that not additional imports are not affecting imports container."""
+    new_parser = OpenAPIParser(
+        source='',
+    )
+    assert len(new_parser.imports) == 0