From b1c0f22ad8c192a3de34a0be4c7671d83cd52b6d Mon Sep 17 00:00:00 2001
From: luca-knaack-webcom <luca.knaack@web-computing.de>
Date: Mon, 10 Jun 2024 21:19:30 +0200
Subject: [PATCH] Add support for external referenced discriminators (#1991)

Co-authored-by: Koudai Aono <koxudaxi@gmail.com>
---
 datamodel_code_generator/parser/base.py       |  6 ++-
 .../output.py                                 | 26 +++++++++++
 .../artificial_folder/type1.json              | 11 +++++
 .../schema.json                               | 44 +++++++++++++++++++
 .../type2.json                                | 11 +++++
 tests/test_main.py                            | 29 ++++++++++++
 6 files changed, 125 insertions(+), 2 deletions(-)
 create mode 100644 tests/data/expected/main/discriminator_with_external_reference/output.py
 create mode 100644 tests/data/jsonschema/discriminator_with_external_reference/artificial_folder/type1.json
 create mode 100644 tests/data/jsonschema/discriminator_with_external_reference/schema.json
 create mode 100644 tests/data/jsonschema/discriminator_with_external_reference/type2.json

diff --git a/datamodel_code_generator/parser/base.py b/datamodel_code_generator/parser/base.py
index 3961208b6..6d80891d2 100644
--- a/datamodel_code_generator/parser/base.py
+++ b/datamodel_code_generator/parser/base.py
@@ -778,8 +778,10 @@ def __apply_discriminator_type(
                                 discriminator_model.path.split('#/')[-1]
                                 != path.split('#/')[-1]
                             ):
-                                # TODO: support external reference
-                                continue
+                                if '#' in path or discriminator_model.path[
+                                    :-1
+                                ] != path.lstrip('./'):
+                                    continue
                             type_names.append(name)
                     else:
                         type_names = [discriminator_model.path.split('/')[-1]]
diff --git a/tests/data/expected/main/discriminator_with_external_reference/output.py b/tests/data/expected/main/discriminator_with_external_reference/output.py
new file mode 100644
index 000000000..a2a1f1d8a
--- /dev/null
+++ b/tests/data/expected/main/discriminator_with_external_reference/output.py
@@ -0,0 +1,26 @@
+# generated by datamodel-codegen:
+#   filename:  schema.json
+#   timestamp: 2019-07-26T00:00:00+00:00
+
+from __future__ import annotations
+
+from typing import Union
+
+from pydantic import BaseModel, Field
+from typing_extensions import Literal
+
+
+class Type1(BaseModel):
+    type_: Literal['a'] = Field('a', title='Type ')
+
+
+class Type2(BaseModel):
+    type_: Literal['b'] = Field('b', title='Type ')
+
+
+class Type3(BaseModel):
+    type_: Literal['c'] = Field('c', title='Type ')
+
+
+class Response(BaseModel):
+    inner: Union[Type1, Type2, Type3] = Field(..., discriminator='type_', title='Inner')
diff --git a/tests/data/jsonschema/discriminator_with_external_reference/artificial_folder/type1.json b/tests/data/jsonschema/discriminator_with_external_reference/artificial_folder/type1.json
new file mode 100644
index 000000000..e7da9f40c
--- /dev/null
+++ b/tests/data/jsonschema/discriminator_with_external_reference/artificial_folder/type1.json
@@ -0,0 +1,11 @@
+{
+  "properties": {
+    "type_": {
+      "const": "a",
+      "default": "a",
+      "title": "Type "
+    }
+  },
+  "title": "Type1",
+  "type": "object"
+}
\ No newline at end of file
diff --git a/tests/data/jsonschema/discriminator_with_external_reference/schema.json b/tests/data/jsonschema/discriminator_with_external_reference/schema.json
new file mode 100644
index 000000000..0fb5c52c7
--- /dev/null
+++ b/tests/data/jsonschema/discriminator_with_external_reference/schema.json
@@ -0,0 +1,44 @@
+{
+  "$def": {
+    "Type3": {
+      "properties": {
+        "type_": {
+          "const": "c",
+          "default": "c",
+          "title": "Type "
+        }
+      },
+      "title": "Type3",
+      "type": "object"
+    }
+  },
+  "properties": {
+    "inner": {
+      "discriminator": {
+        "mapping": {
+          "a": "./artificial_folder/type1.json",
+          "b": "./type2.json",
+          "c": "#/$def/Type3"
+        },
+        "propertyName": "type_"
+      },
+      "oneOf": [
+        {
+          "$ref": "./artificial_folder/type1.json"
+        },
+        {
+          "$ref": "./type2.json"
+        },
+        {
+          "$ref": "#/$def/Type3"
+        }
+      ],
+      "title": "Inner"
+    }
+  },
+  "required": [
+    "inner"
+  ],
+  "title": "Response",
+  "type": "object"
+}
diff --git a/tests/data/jsonschema/discriminator_with_external_reference/type2.json b/tests/data/jsonschema/discriminator_with_external_reference/type2.json
new file mode 100644
index 000000000..aab0950a1
--- /dev/null
+++ b/tests/data/jsonschema/discriminator_with_external_reference/type2.json
@@ -0,0 +1,11 @@
+{
+  "properties": {
+    "type_": {
+      "const": "b",
+      "default": "b",
+      "title": "Type "
+    }
+  },
+  "title": "Type2",
+  "type": "object"
+}
\ No newline at end of file
diff --git a/tests/test_main.py b/tests/test_main.py
index 3854bc591..dae1bfcca 100644
--- a/tests/test_main.py
+++ b/tests/test_main.py
@@ -6123,6 +6123,35 @@ def test_main_jsonschema_discriminator_literals():
         )
 
 
+@freeze_time('2019-07-26')
+def test_main_jsonschema_external_discriminator():
+    with TemporaryDirectory() as output_dir:
+        output_file: Path = Path(output_dir) / 'output.py'
+        return_code: Exit = main(
+            [
+                '--input',
+                str(
+                    JSON_SCHEMA_DATA_PATH
+                    / 'discriminator_with_external_reference'
+                    / 'schema.json'
+                ),
+                '--output',
+                str(output_file),
+                '--output-model-type',
+                'pydantic_v2.BaseModel',
+            ]
+        )
+        assert return_code == Exit.OK
+        assert (
+            output_file.read_text()
+            == (
+                EXPECTED_MAIN_PATH
+                / 'discriminator_with_external_reference'
+                / 'output.py'
+            ).read_text()
+        )
+
+
 @freeze_time('2019-07-26')
 @pytest.mark.skipif(
     black.__version__.split('.')[0] == '19',