Allow WithJsonSchema to inject $refs w/ http or https links (p…

…ydantic#9863) Co-authored-by: Sydney Runkle <[email protected]> Co-authored-by: sydney-runkle <[email protected]>
alexeyt101 · Aug 16, 2024 · 823d51f · 823d51f
1 parent 929543c
commit 823d51f
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 13 deletions.
diff --git 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
@@ -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 == {}