From 3fea8332ed300c5439acde2f92e8c66f2ff5b42b Mon Sep 17 00:00:00 2001
From: Sydney Runkle <54324534+sydney-runkle@users.noreply.github.com>
Date: Wed, 15 Nov 2023 13:46:01 -0600
Subject: [PATCH] Fix bug re `custom_init` on members of `Union` (#1076)

---
 src/validators/model.rs             |  5 +++-
 tests/validators/test_model_init.py | 45 +++++++++++++++++++++++++++++
 2 files changed, 49 insertions(+), 1 deletion(-)

diff --git a/src/validators/model.rs b/src/validators/model.rs
index 0299ce5d8..a571ccd9e 100644
--- a/src/validators/model.rs
+++ b/src/validators/model.rs
@@ -257,7 +257,10 @@ impl ModelValidator {
             // this work with from_attributes, and would essentially allow you to
             // handle init vars by adding them to the __init__ signature.
             if let Some(kwargs) = input.as_kwargs(py) {
-                return Ok(self.class.call(py, (), Some(kwargs))?);
+                return self
+                    .class
+                    .call(py, (), Some(kwargs))
+                    .map_err(|e| convert_err(py, e, input));
             }
         }
 
diff --git a/tests/validators/test_model_init.py b/tests/validators/test_model_init.py
index 5521f8da4..b0c28dc86 100644
--- a/tests/validators/test_model_init.py
+++ b/tests/validators/test_model_init.py
@@ -479,3 +479,48 @@ def _wrap_validator(cls, v, validator, info):
     gc.collect()
 
     assert ref() is None
+
+
+def test_model_custom_init_with_union() -> None:
+    class A:
+        def __init__(self, **kwargs):
+            assert 'a' in kwargs
+            self.a = kwargs.get('a')
+
+    class B:
+        def __init__(self, **kwargs):
+            assert 'b' in kwargs
+            self.b = kwargs.get('b')
+
+    schema = {
+        'type': 'union',
+        'choices': [
+            {
+                'type': 'model',
+                'cls': A,
+                'schema': {
+                    'type': 'model-fields',
+                    'fields': {'a': {'type': 'model-field', 'schema': {'type': 'bool'}}},
+                    'model_name': 'A',
+                },
+                'custom_init': True,
+                'ref': '__main__.A:4947206928',
+            },
+            {
+                'type': 'model',
+                'cls': B,
+                'schema': {
+                    'type': 'model-fields',
+                    'fields': {'b': {'type': 'model-field', 'schema': {'type': 'bool'}}},
+                    'model_name': 'B',
+                },
+                'custom_init': True,
+                'ref': '__main__.B:4679932848',
+            },
+        ],
+    }
+
+    validator = SchemaValidator(schema)
+
+    assert validator.validate_python({'a': False}).a is False
+    assert validator.validate_python({'b': True}).b is True