From b6ac8f9fc57ea9f83fe7202094f121ba5c42473b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=89loi=20Rivard?= <eloi@yaal.coop>
Date: Wed, 26 Jun 2024 11:18:40 +0200
Subject: [PATCH] fix: models with multiples extensions can be defined using
 Union

---
 doc/changelog.rst                |  1 +
 doc/tutorial.rst                 |  2 +-
 scim2_models/rfc7643/resource.py |  3 +++
 tests/test_resource_extension.py | 20 ++++++++++++++++++++
 4 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/doc/changelog.rst b/doc/changelog.rst
index fb1e8ea..95a0a9b 100644
--- a/doc/changelog.rst
+++ b/doc/changelog.rst
@@ -12,6 +12,7 @@ Changed
 Fix
 ^^^
 - Allow relative URLs in :data:`~scim2_models.Reference`.
+- Models with multiples extensions could not be initialized. #37
 
 [0.1.7] - 2024-06-16
 --------------------
diff --git a/doc/tutorial.rst b/doc/tutorial.rst
index 38a6fe2..308b5bd 100644
--- a/doc/tutorial.rst
+++ b/doc/tutorial.rst
@@ -206,7 +206,7 @@ Schema extensions
 =================
 
 :rfc:`RFC7643 §3.3 <7643#section-3.3>` extensions are supported.
-Extensions must be passed as resource type parameter, e.g. ``user = User[EnterpriseUser]`` or ``user = User[EnterpriseUser, SuperHero]``.
+Extensions must be passed as resource type parameter, e.g. ``user = User[EnterpriseUser]`` or ``user = User[Union[EnterpriseUser, SuperHero]]``.
 Extensions attributes are accessed with brackets, e.g. ``user[EnterpriseUser].employee_number``.
 
 .. code-block:: python
diff --git a/scim2_models/rfc7643/resource.py b/scim2_models/rfc7643/resource.py
index ace26c8..08890c3 100644
--- a/scim2_models/rfc7643/resource.py
+++ b/scim2_models/rfc7643/resource.py
@@ -134,6 +134,9 @@ def get_extension_models(cls) -> Dict[str, Type]:
         schemas."""
 
         extension_models = cls.__pydantic_generic_metadata__.get("args", [])
+        if len(extension_models) == 1 and get_origin(extension_models[0]) == Union:
+            extension_models = get_args(extension_models[0])
+
         by_schema = {
             ext.model_fields["schemas"].default[0]: ext for ext in extension_models
         }
diff --git a/tests/test_resource_extension.py b/tests/test_resource_extension.py
index 0dc3c04..0533a2e 100644
--- a/tests/test_resource_extension.py
+++ b/tests/test_resource_extension.py
@@ -1,10 +1,14 @@
 import datetime
+from typing import List
+from typing import Optional
+from typing import Union
 
 import pytest
 
 from scim2_models import EnterpriseUser
 from scim2_models import Manager
 from scim2_models import Meta
+from scim2_models import Resource
 from scim2_models import User
 
 
@@ -180,3 +184,19 @@ def test_invalid_setitem():
 
     with pytest.raises(KeyError):
         user[object] = "foobar"
+
+
+class SuperHero(Resource):
+    schemas: List[str] = ["urn:ietf:params:scim:schemas:extension:enterprise:2.0:User"]
+
+    superpower: Optional[str] = None
+    """The superhero superpower."""
+
+
+def test_multiple_extensions_union():
+    """Test that multiple extensions can be used by using Union."""
+
+    user_model = User[Union[EnterpriseUser, SuperHero]]
+    instance = user_model()
+    instance[SuperHero] = SuperHero(superpower="flight")
+    assert instance[SuperHero].superpower == "flight"