Replace dict.get and dict.setdefault with more verbose versions in Co…

…reSchema building hot paths (#7536)
pydantic · Sep 21, 2023 · 198c8c6 · 198c8c6
1 parent ea7bf54
commit 198c8c6
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 12 deletions.
diff --git a/pydantic/_internal/_core_utils.py b/pydantic/_internal/_core_utils.py
@@ -558,7 +558,10 @@ def inline_refs(s: core_schema.CoreSchema, recurse: Recurse) -> core_schema.Core
 
     if definitions:
         schema = core_schema.definitions_schema(schema=schema, definitions=definitions)
-    schema.setdefault('metadata', {})[_DEFINITIONS_CACHE_METADATA_KEY] = definitions_cache  # type: ignore
+    if 'metadata' in schema:
+        schema['metadata'][_DEFINITIONS_CACHE_METADATA_KEY] = definitions_cache
+    else:
+        schema['metadata'] = {_DEFINITIONS_CACHE_METADATA_KEY: definitions_cache}
     return schema
 
 

diff --git a/pydantic/_internal/_discriminated_union.py b/pydantic/_internal/_discriminated_union.py
@@ -41,12 +41,13 @@ def inner(s: core_schema.CoreSchema, recurse: _core_utils.Recurse) -> core_schem
         if s['type'] == 'tagged-union':
             return s
 
-        metadata = s.get('metadata', {})
-        discriminator = metadata.get(CORE_SCHEMA_METADATA_DISCRIMINATOR_PLACEHOLDER_KEY, None)
-        if discriminator is not None:
-            if definitions is None:
-                definitions = collect_definitions(schema)
-            s = apply_discriminator(s, discriminator, definitions)
+        metadata = s.get('metadata', None)
+        if metadata is not None:
+            discriminator = metadata.get(CORE_SCHEMA_METADATA_DISCRIMINATOR_PLACEHOLDER_KEY, None)
+            if discriminator is not None:
+                if definitions is None:
+                    definitions = collect_definitions(schema)
+                s = apply_discriminator(s, discriminator, definitions)
         return s
 
     return _core_utils.walk_core_schema(schema, inner)

diff --git a/pydantic/_internal/_generate_schema.py b/pydantic/_internal/_generate_schema.py
@@ -370,7 +370,10 @@ def _apply_discriminator_to_union(self, schema: CoreSchema, discriminator: Any)
                 schema,
                 discriminator,
             )
-            schema.setdefault('metadata', {})[NEEDS_APPLY_DISCRIMINATED_UNION_METADATA_KEY] = True
+            if 'metadata' in schema:
+                schema['metadata'][NEEDS_APPLY_DISCRIMINATED_UNION_METADATA_KEY] = True
+            else:
+                schema['metadata'] = {NEEDS_APPLY_DISCRIMINATED_UNION_METADATA_KEY: True}
             self._needs_apply_discriminated_union = True
             return schema
 
@@ -644,9 +647,10 @@ def _get_first_two_args_or_any(self, obj: Any) -> tuple[Any, Any]:
         return args[0], args[1]
 
     def _post_process_generated_schema(self, schema: core_schema.CoreSchema) -> core_schema.CoreSchema:
-        schema.setdefault('metadata', {})[
-            NEEDS_APPLY_DISCRIMINATED_UNION_METADATA_KEY
-        ] = self._needs_apply_discriminated_union
+        if 'metadata' in schema:
+            schema['metadata'][NEEDS_APPLY_DISCRIMINATED_UNION_METADATA_KEY] = self._needs_apply_discriminated_union
+        else:
+            schema['metadata'] = {NEEDS_APPLY_DISCRIMINATED_UNION_METADATA_KEY: self._needs_apply_discriminated_union}
         return schema
 
     def _generate_schema(self, obj: Any) -> core_schema.CoreSchema:

diff --git a/pydantic/_internal/_known_annotated_metadata.py b/pydantic/_internal/_known_annotated_metadata.py
@@ -93,7 +93,14 @@ def update_js_schema(s: cs.CoreSchema, handler: GetJsonSchemaHandler) -> dict[st
         js_schema.update(f())
         return js_schema
 
-    s.setdefault('metadata', {}).setdefault('pydantic_js_functions', []).append(update_js_schema)
+    if 'metadata' in s:
+        metadata = s['metadata']
+        if 'pydantic_js_functions' in s:
+            metadata['pydantic_js_functions'].append(update_js_schema)
+        else:
+            metadata['pydantic_js_functions'] = [update_js_schema]
+    else:
+        s['metadata'] = {'pydantic_js_functions': [update_js_schema]}
 
 
 def as_jsonable_value(v: Any) -> Any: