refactor: extract generic metaclass logic from ListResponse

Author: azmeuk

scim2_models/rfc7644/list_response.py

@@ -2,84 +2,22 @@
 from typing import Any
 from typing import Generic
 from typing import Optional
-from typing import Union
-from typing import get_args
-from typing import get_origin
 
-from pydantic import Discriminator
 from pydantic import Field
-from pydantic import Tag
 from pydantic import ValidationInfo
 from pydantic import ValidatorFunctionWrapHandler
 from pydantic import model_validator
 from pydantic_core import PydanticCustomError
 from typing_extensions import Self
 
-from ..base import BaseModel
-from ..base import BaseModelType
 from ..base import Context
 from ..base import Required
 from ..rfc7643.resource import AnyResource
-from ..utils import UNION_TYPES
+from .message import GenericMessageMetaclass
 from .message import Message
 
 
-class ListResponseMetaclass(BaseModelType):
-    def tagged_resource_union(resource_union):
-        """Build Discriminated Unions, so pydantic can guess which class are needed to instantiate by inspecting a payload.
-
-        https://docs.pydantic.dev/latest/concepts/unions/#discriminated-unions
-        """
-        if get_origin(resource_union) not in UNION_TYPES:
-            return resource_union
-
-        resource_types = get_args(resource_union)
-
-        def get_schema_from_payload(payload: Any) -> Optional[str]:
-            if not payload:
-                return None
-
-            payload_schemas = (
-                payload.get("schemas", [])
-                if isinstance(payload, dict)
-                else payload.schemas
-            )
-
-            resource_types_schemas = [
-                resource_type.model_fields["schemas"].default[0]
-                for resource_type in resource_types
-            ]
-            common_schemas = [
-                schema for schema in payload_schemas if schema in resource_types_schemas
-            ]
-            return common_schemas[0] if common_schemas else None
-
-        discriminator = Discriminator(get_schema_from_payload)
-
-        def get_tag(resource_type: type[BaseModel]) -> Tag:
-            return Tag(resource_type.model_fields["schemas"].default[0])
-
-        tagged_resources = [
-            Annotated[resource_type, get_tag(resource_type)]
-            for resource_type in resource_types
-        ]
-        union = Union[tuple(tagged_resources)]
-        return Annotated[union, discriminator]
-
-    def __new__(cls, name, bases, attrs, **kwargs):
-        if kwargs.get("__pydantic_generic_metadata__") and kwargs[
-            "__pydantic_generic_metadata__"
-        ].get("args"):
-            tagged_union = cls.tagged_resource_union(
-                kwargs["__pydantic_generic_metadata__"]["args"][0]
-            )
-            kwargs["__pydantic_generic_metadata__"]["args"] = (tagged_union,)
-
-        klass = super().__new__(cls, name, bases, attrs, **kwargs)
-        return klass
-
-
-class ListResponse(Message, Generic[AnyResource], metaclass=ListResponseMetaclass):
+class ListResponse(Message, Generic[AnyResource], metaclass=GenericMessageMetaclass):
     schemas: Annotated[list[str], Required.true] = [
         "urn:ietf:params:scim:api:messages:2.0:ListResponse"
     ]

scim2_models/rfc7644/message.py

@@ -1,10 +1,98 @@
 from typing import Annotated
+from typing import Any
+from typing import Optional
+from typing import Union
+from typing import get_args
+from typing import get_origin
+
+from pydantic import Discriminator
+from pydantic import Tag
 
 from ..base import BaseModel
+from ..base import BaseModelType
 from ..base import Required
+from ..utils import UNION_TYPES
 
 
 class Message(BaseModel):
     """SCIM protocol messages as defined by :rfc:`RFC7644 §3.1 <7644#section-3.1>`."""
 
     schemas: Annotated[list[str], Required.true]
+
+
+def get_schema_from_payload(payload: Any) -> Optional[str]:
+    """Extract schema from SCIM payload for discrimination.
+
+    :param payload: SCIM payload dict or object
+    :return: First matching schema or None
+    """
+    if not payload:
+        return None
+
+    payload_schemas = (
+        payload.get("schemas", []) if isinstance(payload, dict) else payload.schemas
+    )
+
+    # This will be set by the calling context
+    resource_types_schemas = getattr(get_schema_from_payload, "_resource_schemas", [])
+    common_schemas = [
+        schema for schema in payload_schemas if schema in resource_types_schemas
+    ]
+    return common_schemas[0] if common_schemas else None
+
+
+def get_tag(resource_type: type[BaseModel]) -> Tag:
+    """Create Pydantic tag from resource type schema.
+
+    :param resource_type: SCIM resource type
+    :return: Pydantic Tag for discrimination
+    """
+    return Tag(resource_type.model_fields["schemas"].default[0])
+
+
+def create_tagged_resource_union(resource_union):
+    """Build Discriminated Unions for SCIM resources.
+
+    Creates discriminated unions so Pydantic can determine which class to instantiate
+    by inspecting the payload's schemas field.
+
+    :param resource_union: Union type of SCIM resources
+    :return: Annotated discriminated union or original type
+    """
+    if get_origin(resource_union) not in UNION_TYPES:
+        return resource_union
+
+    resource_types = get_args(resource_union)
+
+    # Set up schemas for the discriminator function
+    resource_types_schemas = [
+        resource_type.model_fields["schemas"].default[0]
+        for resource_type in resource_types
+    ]
+    get_schema_from_payload._resource_schemas = resource_types_schemas
+
+    discriminator = Discriminator(get_schema_from_payload)
+
+    tagged_resources = [
+        Annotated[resource_type, get_tag(resource_type)]
+        for resource_type in resource_types
+    ]
+    union = Union[tuple(tagged_resources)]
+    return Annotated[union, discriminator]
+
+
+class GenericMessageMetaclass(BaseModelType):
+    """Metaclass for SCIM generic types with discriminated unions."""
+
+    def __new__(cls, name, bases, attrs, **kwargs):
+        """Create class with tagged resource unions for generic parameters."""
+        if kwargs.get("__pydantic_generic_metadata__") and kwargs[
+            "__pydantic_generic_metadata__"
+        ].get("args"):
+            tagged_union = create_tagged_resource_union(
+                kwargs["__pydantic_generic_metadata__"]["args"][0]
+            )
+            kwargs["__pydantic_generic_metadata__"]["args"] = (tagged_union,)
+
+        klass = super().__new__(cls, name, bases, attrs, **kwargs)
+        return klass