From 823d51f92f544614bca556076e27dbacd1298491 Mon Sep 17 00:00:00 2001 From: dAIsySHEng1 <45747761+dAIsySHEng1@users.noreply.github.com> Date: Fri, 16 Aug 2024 12:18:05 -0400 Subject: [PATCH] Allow `WithJsonSchema` to inject `$ref`s w/ `http` or `https` links (#9863) Co-authored-by: Sydney Runkle <54324534+sydney-runkle@users.noreply.github.com> Co-authored-by: sydney-runkle --- pydantic/json_schema.py | 39 ++++++++++++++++++++++++++------------- tests/test_json_schema.py | 17 +++++++++++++++++ 2 files changed, 43 insertions(+), 13 deletions(-) diff --git a/pydantic/json_schema.py b/pydantic/json_schema.py index dd99c3f0f7..f7f3b3289d 100644 --- a/pydantic/json_schema.py +++ b/pydantic/json_schema.py @@ -2050,10 +2050,15 @@ def handle_ref_overrides(self, json_schema: JsonSchemaValue) -> JsonSchemaValue: return json_schema def get_schema_from_definitions(self, json_ref: JsonRef) -> JsonSchemaValue | None: - def_ref = self.json_to_defs_refs[json_ref] - if def_ref in self._core_defs_invalid_for_json_schema: - raise self._core_defs_invalid_for_json_schema[def_ref] - return self.definitions.get(def_ref, None) + try: + def_ref = self.json_to_defs_refs[json_ref] + if def_ref in self._core_defs_invalid_for_json_schema: + raise self._core_defs_invalid_for_json_schema[def_ref] + return self.definitions.get(def_ref, None) + except KeyError: + if json_ref.startswith(('http://', 'https://')): + return None + raise def encode_default(self, dft: Any) -> Any: """Encode a default value to a JSON-serializable value. @@ -2163,10 +2168,14 @@ def _add_json_refs(schema: Any) -> None: json_refs[json_ref] += 1 if already_visited: return # prevent recursion on a definition that was already visited - defs_ref = self.json_to_defs_refs[json_ref] - if defs_ref in self._core_defs_invalid_for_json_schema: - raise self._core_defs_invalid_for_json_schema[defs_ref] - _add_json_refs(self.definitions[defs_ref]) + try: + defs_ref = self.json_to_defs_refs[json_ref] + if defs_ref in self._core_defs_invalid_for_json_schema: + raise self._core_defs_invalid_for_json_schema[defs_ref] + _add_json_refs(self.definitions[defs_ref]) + except KeyError: + if not json_ref.startswith(('http://', 'https://')): + raise for v in schema.values(): _add_json_refs(v) @@ -2223,11 +2232,15 @@ def _garbage_collect_definitions(self, schema: JsonSchemaValue) -> None: unvisited_json_refs = _get_all_json_refs(schema) while unvisited_json_refs: next_json_ref = unvisited_json_refs.pop() - next_defs_ref = self.json_to_defs_refs[next_json_ref] - if next_defs_ref in visited_defs_refs: - continue - visited_defs_refs.add(next_defs_ref) - unvisited_json_refs.update(_get_all_json_refs(self.definitions[next_defs_ref])) + try: + next_defs_ref = self.json_to_defs_refs[next_json_ref] + if next_defs_ref in visited_defs_refs: + continue + visited_defs_refs.add(next_defs_ref) + unvisited_json_refs.update(_get_all_json_refs(self.definitions[next_defs_ref])) + except KeyError: + if not next_json_ref.startswith(('http://', 'https://')): + raise self.definitions = {k: v for k, v in self.definitions.items() if k in visited_defs_refs} diff --git a/tests/test_json_schema.py b/tests/test_json_schema.py index c9316c99a4..6adb69e93f 100644 --- a/tests/test_json_schema.py +++ b/tests/test_json_schema.py @@ -1664,6 +1664,23 @@ class Baz(BaseModel): models_json_schema([(Bar, 'validation'), (Baz, 'validation')], ref_template='/schemas/{bad_name}.json#/') +def test_external_ref(): + """https://github.com/pydantic/pydantic/issues/9783""" + + class Model(BaseModel): + json_schema: Annotated[ + dict, + WithJsonSchema({'$ref': 'https://json-schema.org/draft/2020-12/schema'}), + ] + + assert Model.model_json_schema() == { + 'properties': {'json_schema': {'$ref': 'https://json-schema.org/draft/2020-12/schema', 'title': 'Json Schema'}}, + 'required': ['json_schema'], + 'title': 'Model', + 'type': 'object', + } + + def test_schema_no_definitions(): keys_map, model_schema = models_json_schema([], title='Schema without definitions') assert keys_map == {}