refactor: move get_reference_types in reference.py

Author: azmeuk

scim2_models/reference.py

@@ -2,11 +2,15 @@
 from typing import Any
 from typing import Generic
 from typing import TypeVar
+from typing import get_args
+from typing import get_origin
 
 from pydantic import GetCoreSchemaHandler
 from pydantic_core import core_schema
 from typing_extensions import NewType
 
+from .utils import UNION_TYPES
+
 ReferenceTypes = TypeVar("ReferenceTypes")
 
 URIReference = NewType("URIReference", str)
@@ -48,3 +52,28 @@ def __get_pydantic_core_schema__(
     @classmethod
     def _validate(cls, input_value: str, /) -> str:
         return input_value
+
+    @classmethod
+    def get_types(cls, type_annotation: Any) -> list[str]:
+        """Get reference types from a type annotation.
+
+        :param type_annotation: Type annotation to extract reference types from
+        :type type_annotation: Any
+        :return: List of reference type strings
+        :rtype: list[str]
+        """
+        first_arg = get_args(type_annotation)[0]
+        types = (
+            get_args(first_arg) if get_origin(first_arg) in UNION_TYPES else [first_arg]
+        )
+
+        def serialize_ref_type(ref_type: Any) -> str:
+            if ref_type == URIReference:
+                return "uri"
+
+            elif ref_type == ExternalReference:
+                return "external"
+
+            return get_args(ref_type)[0]
+
+        return list(map(serialize_ref_type, types))

scim2_models/rfc7643/resource.py

@@ -29,8 +29,7 @@
 from ..attributes import is_complex_attribute
 from ..base import BaseModel
 from ..base import BaseModelType
-from ..reference import ExternalReference
-from ..reference import URIReference
+from ..reference import Reference
 from ..utils import UNION_TYPES
 from ..utils import normalize_attribute_name
 
@@ -317,22 +316,6 @@ def model_to_schema(model: type[BaseModel]) -> "Schema":
     return schema
 
 
-def get_reference_types(type) -> list[str]:
-    first_arg = get_args(type)[0]
-    types = get_args(first_arg) if get_origin(first_arg) in UNION_TYPES else [first_arg]
-
-    def serialize_ref_type(ref_type: type) -> str:
-        if ref_type == URIReference:
-            return "uri"
-
-        elif ref_type == ExternalReference:
-            return "external"
-
-        return get_args(ref_type)[0]
-
-    return list(map(serialize_ref_type, types))
-
-
 def model_attribute_to_scim_attribute(
     model: type[BaseModel], attribute_name: str
 ) -> "Attribute":
@@ -369,7 +352,7 @@ def model_attribute_to_scim_attribute(
         returned=model.get_field_annotation(attribute_name, Returned),
         uniqueness=model.get_field_annotation(attribute_name, Uniqueness),
         sub_attributes=sub_attributes,
-        reference_types=get_reference_types(root_type)
+        reference_types=Reference.get_types(root_type)
         if attribute_type == Attribute.Type.reference
         else None,
     )