fix: components schema key names (#2841)

* fix: components schema key names

219355c#diff-6cdcc2ab6bdeda3d72588f3b65db0c437f7ba9cad43892829e75c4d701f9e658R659-R661 introduced a change to the way that keys for the openapi components/schemas object were calculated to address name the possibility of name collisions.

This PR reverts that behavior for the case where a name has no collision, and only introduces extended keys for the case where there are multiple objects with the same name.

Closes #2804

* refactor: replace schemas with schema_registry

* docs: docstrings for things

* refactor: removes OpenAPIContext.schemas

In favor of OpenAPIContext.schema_registry

* refactor: remove unused slots

* refactor: modify the schema key generation function

This makes it appropriate as a function that generates a key for the registry, not a key for components/schemas.

* Update tests/unit/test_contrib/test_repository.py

* fix: don't add arrays to components/schemas

litestar/_openapi/datastructures.py

@@ -1,39 +1,152 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Sequence
+from collections import defaultdict
+from typing import TYPE_CHECKING, Iterator, Sequence
 
 from litestar.exceptions import ImproperlyConfiguredException
 
 if TYPE_CHECKING:
     from litestar.openapi import OpenAPIConfig
-    from litestar.openapi.spec import Schema
+    from litestar.openapi.spec import Reference, Schema
     from litestar.plugins import OpenAPISchemaPluginProtocol
 
 
-class OpenAPIContext:
-    """OpenAPI Context.
+class RegisteredSchema:
+    """Object to store a schema and any references to it."""
+
+    def __init__(self, key: tuple[str, ...], schema: Schema, references: list[Reference]) -> None:
+        """Create a new RegisteredSchema object.
+
+        Args:
+            key: The key used to register the schema.
+            schema: The schema object.
+            references: A list of references to the schema.
+        """
+        self.key = key
+        self.schema = schema
+        self.references = references
+
+
+class SchemaRegistry:
+    """A registry for object schemas.
+
+    This class is used to store schemas that we reference from other parts of the spec.
+
+    Its main purpose is to allow us to generate the components/schemas section of the spec once we have
+    collected all the schemas that should be included.
 
-    Context object used to support OpenAPI schema generation.
+    This allows us to determine a path to the schema in the components/schemas section of the spec that
+    is unique and as short as possible.
     """
 
-    __slots__ = ("openapi_config", "plugins", "schemas", "operation_ids")
+    def __init__(self) -> None:
+        self._schema_key_map: dict[tuple[str, ...], RegisteredSchema] = {}
+        self._schema_reference_map: dict[int, RegisteredSchema] = {}
+        self._model_name_groups: defaultdict[str, list[RegisteredSchema]] = defaultdict(list)
 
-    def __init__(
-        self, openapi_config: OpenAPIConfig, plugins: Sequence[OpenAPISchemaPluginProtocol], schemas: dict[str, Schema]
+    def register(
+        self,
+        key: tuple[str, ...],
+        schema: Schema,
+        reference: Reference,
     ) -> None:
-        """Initialize OpenAPIContext.
+        if (registered_schema := self._schema_key_map.get(key)) is not None:
+            registered_schema.references.append(reference)
+            self._schema_reference_map[id(reference)] = registered_schema
+            return
+
+        self._schema_key_map[key] = registered_schema = RegisteredSchema(key, schema, [reference])
+        self._schema_reference_map[id(reference)] = registered_schema
+        self._model_name_groups[key[-1]].append(registered_schema)
+
+    def from_reference(self, reference: Reference) -> RegisteredSchema:
+        """Get a registered schema by its reference.
 
         Args:
-            openapi_config: OpenAPIConfig instance.
-            plugins: OpenAPI plugins.
-            schemas: Mapping of schema names to schema objects that will become the components.schemas section of the
-                OpenAPI schema.
+            reference: The reference to the schema to get.
+
+        Returns:
+            A RegisteredSchema object.
         """
+        return self._schema_reference_map[id(reference)]
+
+    def __iter__(self) -> Iterator[RegisteredSchema]:
+        """Iterate over the registered schemas."""
+        return iter(self._schema_key_map.values())
+
+    @staticmethod
+    def set_reference_paths(name: str, registered_schema: RegisteredSchema) -> None:
+        """Set the reference paths for a registered schema."""
+        for reference in registered_schema.references:
+            reference.ref = f"#/components/schemas/{name}"
+
+    @staticmethod
+    def remove_common_prefix(tuples: list[tuple[str, ...]]) -> list[tuple[str, ...]]:
+        """Remove the common prefix from a list of tuples.
+
+        Args:
+            tuples: A list of tuples to remove the common prefix from.
+
+        Returns:
+            A list of tuples with the common prefix removed.
+        """
+
+        def longest_common_prefix(tuples_: list[tuple[str, ...]]) -> tuple[str, ...]:
+            """Find the longest common prefix of a list of tuples.
+
+            Args:
+                tuples_: A list of tuples to find the longest common prefix of.
+
+            Returns:
+                The longest common prefix of the tuples.
+            """
+            prefix_ = tuples_[0]
+            for t in tuples_:
+                # Compare the current prefix with each tuple and shorten it
+                prefix_ = prefix_[: min(len(prefix_), len(t))]
+                for i in range(len(prefix_)):
+                    if prefix_[i] != t[i]:
+                        prefix_ = prefix_[:i]
+                        break
+            return prefix_
+
+        prefix = longest_common_prefix(tuples)
+        prefix_length = len(prefix)
+        return [t[prefix_length:] for t in tuples]
+
+    def generate_components_schemas(self) -> dict[str, Schema]:
+        """Generate the components/schemas section of the spec.
+
+        Returns:
+            A dictionary of schemas.
+        """
+        components_schemas: dict[str, Schema] = {}
+
+        for name, name_group in self._model_name_groups.items():
+            if len(name_group) == 1:
+                self.set_reference_paths(name, name_group[0])
+                components_schemas[name] = name_group[0].schema
+                continue
+
+            full_keys = [registered_schema.key for registered_schema in name_group]
+            names = ["_".join(k) for k in self.remove_common_prefix(full_keys)]
+            for name_, registered_schema in zip(names, name_group):
+                self.set_reference_paths(name_, registered_schema)
+                components_schemas[name_] = registered_schema.schema
+
+        return components_schemas
+
+
+class OpenAPIContext:
+    def __init__(
+        self,
+        openapi_config: OpenAPIConfig,
+        plugins: Sequence[OpenAPISchemaPluginProtocol],
+    ) -> None:
         self.openapi_config = openapi_config
         self.plugins = plugins
-        self.schemas = schemas
-        # used to track that operation ids are globally unique across the OpenAPI document
         self.operation_ids: set[str] = set()
+        self.schema_registry = SchemaRegistry()
 
     def add_operation_id(self, operation_id: str) -> None:
         """Add an operation ID to the context.

litestar/_openapi/parameters.py

@@ -131,7 +131,7 @@ def create_parameter(self, field_definition: FieldDefinition, parameter_name: st
         if not result:
             result = self.schema_creator.for_field_definition(field_definition)
 
-        schema = result if isinstance(result, Schema) else self.context.schemas[result.value]
+        schema = result if isinstance(result, Schema) else self.context.schema_registry.from_reference(result).schema
 
         examples_list = kwarg_definition.examples or [] if kwarg_definition else []
         examples = get_formatted_examples(field_definition, examples_list)

litestar/_openapi/plugin.py

@@ -33,14 +33,11 @@ def __init__(self, app: Litestar) -> None:
 
     def _build_openapi_schema(self) -> OpenAPI:
         openapi = self.openapi_config.to_openapi_schema()
-        context = OpenAPIContext(
-            openapi_config=self.openapi_config,
-            plugins=self.app.plugins.openapi,
-            schemas=openapi.components.schemas,
-        )
+        context = OpenAPIContext(openapi_config=self.openapi_config, plugins=self.app.plugins.openapi)
         openapi.paths = {
             route.path_format or "/": create_path_item_for_route(context, route) for route in self.included_routes
         }
+        openapi.components.schemas = context.schema_registry.generate_components_schemas()
         return openapi
 
     def provide_openapi(self) -> OpenAPI:

litestar/_openapi/responses.py

@@ -148,7 +148,9 @@ def create_success_response(self) -> OpenAPIResponse:
 
                 result = self.schema_creator.for_field_definition(field_def)
 
-            schema = result if isinstance(result, Schema) else self.context.schemas[result.value]
+            schema = (
+                result if isinstance(result, Schema) else self.context.schema_registry.from_reference(result).schema
+            )
             schema.content_encoding = self.route_handler.content_encoding
             schema.content_media_type = self.route_handler.content_media_type
             response = OpenAPIResponse(

litestar/_openapi/schema_generation/schema.py

@@ -37,6 +37,7 @@
 from msgspec.structs import fields as msgspec_struct_fields
 from typing_extensions import NotRequired, Required, Self, get_args
 
+from litestar._openapi.datastructures import SchemaRegistry
 from litestar._openapi.schema_generation.constrained_fields import (
     create_date_constrained_field_schema,
     create_numerical_constrained_field_schema,
@@ -260,41 +261,41 @@ def create_schema_for_annotation(annotation: Any) -> Schema:
 
 
 class SchemaCreator:
-    __slots__ = ("generate_examples", "plugins", "schemas", "prefer_alias", "dto_for")
+    __slots__ = ("generate_examples", "plugins", "prefer_alias", "schema_registry")
 
     def __init__(
         self,
         generate_examples: bool = False,
         plugins: Iterable[OpenAPISchemaPluginProtocol] | None = None,
-        schemas: dict[str, Schema] | None = None,
         prefer_alias: bool = True,
+        schema_registry: SchemaRegistry | None = None,
     ) -> None:
         """Instantiate a SchemaCreator.
 
         Args:
             generate_examples: Whether to generate examples if none are given.
             plugins: A list of plugins.
-            schemas: A mapping of namespaces to schemas - this mapping is used in the OA components section.
             prefer_alias: Whether to prefer the alias name for the schema.
+            schema_registry: A SchemaRegistry instance.
         """
         self.generate_examples = generate_examples
         self.plugins = plugins if plugins is not None else []
-        self.schemas = schemas if schemas is not None else {}
         self.prefer_alias = prefer_alias
+        self.schema_registry = schema_registry or SchemaRegistry()
 
     @classmethod
     def from_openapi_context(cls, context: OpenAPIContext, prefer_alias: bool = True, **kwargs: Any) -> Self:
         kwargs.setdefault("generate_examples", context.openapi_config.create_examples)
         kwargs.setdefault("plugins", context.plugins)
-        kwargs.setdefault("schemas", context.schemas)
+        kwargs.setdefault("schema_registry", context.schema_registry)
         return cls(**kwargs, prefer_alias=prefer_alias)
 
     @property
     def not_generating_examples(self) -> SchemaCreator:
         """Return a SchemaCreator with generate_examples set to False."""
         if not self.generate_examples:
             return self
-        return type(self)(generate_examples=False, plugins=self.plugins, schemas=self.schemas, prefer_alias=False)
+        return type(self)(generate_examples=False, plugins=self.plugins, prefer_alias=False)
 
     def get_plugin_for(self, field_definition: FieldDefinition) -> OpenAPISchemaPluginProtocol | None:
         return next(
@@ -689,12 +690,12 @@ def process_schema_result(self, field: FieldDefinition, schema: Schema) -> Schem
 
             schema.examples = get_formatted_examples(field, create_examples_for_field(field))
 
-        if schema.title and schema.type in (OpenAPIType.OBJECT, OpenAPIType.ARRAY):
-            class_name = _get_normalized_schema_key(str(field.annotation))
-
-            if class_name in self.schemas:
-                return Reference(ref=f"#/components/schemas/{class_name}", description=schema.description)
-
-            self.schemas[class_name] = schema
-            return Reference(ref=f"#/components/schemas/{class_name}")
+        if schema.title and schema.type == OpenAPIType.OBJECT:
+            class_key = _get_normalized_schema_key(field.annotation)
+            # the "ref" attribute set here is arbitrary, since it will be overwritten by the SchemaRegistry
+            # when the "components/schemas" section of the OpenAPI document is generated and the paths are
+            # known.
+            ref = Reference(ref="", description=schema.description)
+            self.schema_registry.register(class_key, schema, ref)
+            return ref
         return schema

litestar/_openapi/schema_generation/utils.py

@@ -1,8 +1,7 @@
 from __future__ import annotations
 
-import re
 from enum import Enum
-from typing import TYPE_CHECKING, Any, Mapping
+from typing import TYPE_CHECKING, Any, Mapping, _GenericAlias  # type: ignore[attr-defined]
 
 from litestar.utils.helpers import get_name
 
@@ -84,33 +83,21 @@ def _should_create_literal_schema(field_definition: FieldDefinition) -> bool:
     )
 
 
-TYPE_NAME_NORMALIZATION_SUB_REGEX = re.compile(r"[^a-zA-Z0-9]+")
-TYPE_NAME_EXTRACTION_REGEX = re.compile(r"<\w+ '(.+)'")
+def _get_normalized_schema_key(annotation: Any) -> tuple[str, ...]:
+    """Create a key for a type annotation.
 
-
-def _replace_non_alphanumeric_match(match: re.Match) -> str:
-    # we don't want to introduce leading or trailing underscores, so we only replace a
-    # char with an underscore if we're not at the beginning or at the end of the
-    # matchable string
-    return "" if match.start() == 0 or match.end() == match.endpos else "_"
-
-
-def _get_normalized_schema_key(type_annotation_str: str) -> str:
-    """Normalize a type annotation, replacing all non-alphanumeric with underscores.
-    Existing underscores will be left as-is
+    The key should be a tuple such as ``("path", "to", "type", "TypeName")``.
 
     Args:
-        type_annotation_str: A string representing a type annotation
-            (i.e. 'typing.Dict[str, typing.Any]' or '<class 'model.Foo'>')
+        annotation: a type annotation
 
     Returns:
-        A normalized version of the input string
+        A tuple of strings.
     """
-    # extract names from repr() style annotations like <class 'foo.bar.Baz'>
-    normalized_name = TYPE_NAME_EXTRACTION_REGEX.sub(r"\g<1>", type_annotation_str)
-    # replace all non-alphanumeric characters with underscores, ensuring no leading or
-    # trailing underscores
-    return TYPE_NAME_NORMALIZATION_SUB_REGEX.sub(_replace_non_alphanumeric_match, normalized_name)
+    module = getattr(annotation, "__module__", "")
+    name = str(annotation)[len(module) + 1 :] if isinstance(annotation, _GenericAlias) else annotation.__qualname__
+    name = name.replace(".<locals>.", ".")
+    return *module.split("."), name
 
 
 def get_formatted_examples(field_definition: FieldDefinition, examples: Sequence[Example]) -> Mapping[str, Example]:

tests/examples/test_openapi.py

@@ -18,13 +18,7 @@ def test_schema_generation() -> None:
                             "200": {
                                 "description": "Request fulfilled, document follows",
                                 "headers": {},
-                                "content": {
-                                    "application/json": {
-                                        "schema": {
-                                            "$ref": "#/components/schemas/docs_examples_openapi_customize_pydantic_model_name_IdModel"
-                                        }
-                                    }
-                                },
+                                "content": {"application/json": {"schema": {"$ref": "#/components/schemas/IdModel"}}},
                             }
                         },
                         "deprecated": False,
@@ -33,7 +27,7 @@ def test_schema_generation() -> None:
             },
             "components": {
                 "schemas": {
-                    "docs_examples_openapi_customize_pydantic_model_name_IdModel": {
+                    "IdModel": {
                         "properties": {"id": {"type": "string", "format": "uuid", "description": "Any UUID string"}},
                         "type": "object",
                         "required": ["id"],

tests/helpers.py

@@ -6,6 +6,11 @@
 from contextlib import AbstractContextManager
 from typing import Any, AsyncContextManager, Awaitable, ContextManager, TypeVar, cast, overload
 
+from litestar._openapi.schema_generation import SchemaCreator
+from litestar.openapi.spec import Schema
+from litestar.plugins import OpenAPISchemaPluginProtocol
+from litestar.typing import FieldDefinition
+
 T = TypeVar("T")
 
 
@@ -51,3 +56,14 @@ def maybe_async_cm(obj: ContextManager[T] | AsyncContextManager[T]) -> AsyncCont
     if isinstance(obj, AbstractContextManager):
         return cast(AsyncContextManager[T], _AsyncContextManagerWrapper(obj))
     return obj
+
+
+def get_schema_for_field_definition(
+    field_definition: FieldDefinition, *, plugins: list[OpenAPISchemaPluginProtocol] | None = None
+) -> Schema:
+    plugins = plugins or []
+    creator = SchemaCreator(plugins=plugins)
+    result = creator.for_field_definition(field_definition)
+    if isinstance(result, Schema):
+        return result
+    return creator.schema_registry.from_reference(result).schema