Fix Pydantic Unions

RELEASE.md

@@ -0,0 +1,3 @@
+Release type: patch
+
+Fixes a bug in the Pydantic conversion code around `Union` values.

strawberry/experimental/pydantic/conversion.py

@@ -3,6 +3,7 @@
 from strawberry.field import StrawberryField
 from strawberry.scalars import is_scalar
 from strawberry.type import StrawberryList, StrawberryOptional, StrawberryType
+from strawberry.union import StrawberryUnion
 
 
 def _convert_from_pydantic_to_strawberry_type(
@@ -16,6 +17,16 @@ def _convert_from_pydantic_to_strawberry_type(
         return _convert_from_pydantic_to_strawberry_type(
             type_.of_type, data_from_model=data, extra=extra
         )
+    if isinstance(type_, StrawberryUnion):
+        for option_type in type_.types:
+            if hasattr(option_type, "_pydantic_type"):
+                source_type = option_type._pydantic_type  # type: ignore
+            else:
+                source_type = cast(type, option_type)
+            if isinstance(data, source_type):
+                return _convert_from_pydantic_to_strawberry_type(
+                    option_type, data_from_model=data, extra=extra
+                )
     if isinstance(type_, StrawberryList):
         items = []
         for index, item in enumerate(data):

strawberry/experimental/pydantic/object_type.py

@@ -144,6 +144,7 @@ def wrap(cls):
         )
 
         model._strawberry_type = cls  # type: ignore
+        cls._pydantic_type = model  # type: ignore
 
         def from_pydantic(instance: Any, extra: Dict[str, Any] = None) -> Any:
             return convert_pydantic_model_to_strawberry_class(

tests/experimental/pydantic/test_conversion.py

@@ -1,4 +1,4 @@
-from typing import List, Optional
+from typing import List, Optional, Union
 
 import pydantic
 
@@ -227,6 +227,120 @@ class User:
     ]
 
 
+def test_can_convert_pydantic_type_to_strawberry_with_union():
+    class BranchA(pydantic.BaseModel):
+        field_a: str
+
+    class BranchB(pydantic.BaseModel):
+        field_b: int
+
+    class User(pydantic.BaseModel):
+        age: int
+        union_field: Union[BranchA, BranchB]
+
+    @strawberry.experimental.pydantic.type(BranchA, fields=["field_a"])
+    class BranchAType:
+        pass
+
+    @strawberry.experimental.pydantic.type(BranchB, fields=["field_b"])
+    class BranchBType:
+        pass
+
+    @strawberry.experimental.pydantic.type(User, fields=["age", "union_field"])
+    class UserType:
+        pass
+
+    origin_user = User(age=1, union_field=BranchA(field_a="abc"))
+    user = UserType.from_pydantic(origin_user)
+
+    assert user.age == 1
+    assert isinstance(user.union_field, BranchAType)
+    assert user.union_field.field_a == "abc"
+
+    origin_user = User(age=1, union_field=BranchB(field_b=123))
+    user = UserType.from_pydantic(origin_user)
+
+    assert user.age == 1
+    assert isinstance(user.union_field, BranchBType)
+    assert user.union_field.field_b == 123
+
+
+def test_can_convert_pydantic_type_to_strawberry_with_union_of_strawberry_types():
+    @strawberry.type
+    class BranchA:
+        field_a: str
+
+    @strawberry.type
+    class BranchB:
+        field_b: int
+
+    class User(pydantic.BaseModel):
+        age: int
+        union_field: Union[BranchA, BranchB]
+
+    @strawberry.experimental.pydantic.type(User, fields=["age", "union_field"])
+    class UserType:
+        pass
+
+    origin_user = User(age=1, union_field=BranchA(field_a="abc"))
+    user = UserType.from_pydantic(origin_user)
+
+    assert user.age == 1
+    assert isinstance(user.union_field, BranchA)
+    assert user.union_field.field_a == "abc"
+
+    origin_user = User(age=1, union_field=BranchB(field_b=123))
+    user = UserType.from_pydantic(origin_user)
+
+    assert user.age == 1
+    assert isinstance(user.union_field, BranchB)
+    assert user.union_field.field_b == 123
+
+
+def test_can_convert_pydantic_type_to_strawberry_with_union_nullable():
+    class BranchA(pydantic.BaseModel):
+        field_a: str
+
+    class BranchB(pydantic.BaseModel):
+        field_b: int
+
+    class User(pydantic.BaseModel):
+        age: int
+        union_field: Union[None, BranchA, BranchB]
+
+    @strawberry.experimental.pydantic.type(BranchA, fields=["field_a"])
+    class BranchAType:
+        pass
+
+    @strawberry.experimental.pydantic.type(BranchB, fields=["field_b"])
+    class BranchBType:
+        pass
+
+    @strawberry.experimental.pydantic.type(User, fields=["age", "union_field"])
+    class UserType:
+        pass
+
+    origin_user = User(age=1, union_field=BranchA(field_a="abc"))
+    user = UserType.from_pydantic(origin_user)
+
+    assert user.age == 1
+    assert isinstance(user.union_field, BranchAType)
+    assert user.union_field.field_a == "abc"
+
+    origin_user = User(age=1, union_field=BranchB(field_b=123))
+    user = UserType.from_pydantic(origin_user)
+
+    assert user.age == 1
+    assert isinstance(user.union_field, BranchBType)
+    assert user.union_field.field_b == 123
+
+    origin_user = User(age=1, union_field=None)
+    user = UserType.from_pydantic(origin_user)
+
+    assert user.age == 1
+    assert user.union_field is None
+
+
 def test_can_convert_pydantic_type_to_strawberry_with_additional_fields():
     class UserModel(pydantic.BaseModel):
         password: Optional[str]