union ser fix for various lookup key types

Author: sydney-runkle

src/lookup_key.rs

@@ -191,34 +191,10 @@ impl LookupKey {
         }
     }
 
-    pub fn py_get_attr<'py, 's>(
-        &'s self,
-        obj: &Bound<'py, PyAny>,
-        kwargs: Option<&Bound<'py, PyDict>>,
-    ) -> ValResult<Option<(&'s LookupPath, Bound<'py, PyAny>)>> {
-        match self._py_get_attr(obj, kwargs) {
-            Ok(v) => Ok(v),
-            Err(err) => {
-                let error = py_err_string(obj.py(), err);
-                Err(ValError::new(
-                    ErrorType::GetAttributeError { error, context: None },
-                    obj,
-                ))
-            }
-        }
-    }
-
-    pub fn _py_get_attr<'py, 's>(
+    pub fn simple_py_get_attr<'py, 's>(
         &'s self,
         obj: &Bound<'py, PyAny>,
-        kwargs: Option<&Bound<'py, PyDict>>,
     ) -> PyResult<Option<(&'s LookupPath, Bound<'py, PyAny>)>> {
-        if let Some(dict) = kwargs {
-            if let Ok(Some(item)) = self.py_get_dict_item(dict) {
-                return Ok(Some(item));
-            }
-        }
-
         match self {
             Self::Simple { py_key, path, .. } => match py_get_attrs(obj, py_key)? {
                 Some(value) => Ok(Some((path, value))),
@@ -260,6 +236,37 @@ impl LookupKey {
         }
     }
 
+    pub fn py_get_attr<'py, 's>(
+        &'s self,
+        obj: &Bound<'py, PyAny>,
+        kwargs: Option<&Bound<'py, PyDict>>,
+    ) -> ValResult<Option<(&'s LookupPath, Bound<'py, PyAny>)>> {
+        match self._py_get_attr(obj, kwargs) {
+            Ok(v) => Ok(v),
+            Err(err) => {
+                let error = py_err_string(obj.py(), err);
+                Err(ValError::new(
+                    ErrorType::GetAttributeError { error, context: None },
+                    obj,
+                ))
+            }
+        }
+    }
+
+    pub fn _py_get_attr<'py, 's>(
+        &'s self,
+        obj: &Bound<'py, PyAny>,
+        kwargs: Option<&Bound<'py, PyDict>>,
+    ) -> PyResult<Option<(&'s LookupPath, Bound<'py, PyAny>)>> {
+        if let Some(dict) = kwargs {
+            if let Ok(Some(item)) = self.py_get_dict_item(dict) {
+                return Ok(Some(item));
+            }
+        }
+
+        self.simple_py_get_attr(obj)
+    }
+
     pub fn json_get<'a, 'data, 's>(
         &'s self,
         dict: &'a JsonObject<'data>,

src/serializers/type_serializers/union.rs

@@ -8,7 +8,6 @@ use std::borrow::Cow;
 use crate::build_tools::py_schema_err;
 use crate::common::union::{Discriminator, SMALL_UNION_THRESHOLD};
 use crate::definitions::DefinitionsBuilder;
-use crate::lookup_key::LookupKey;
 use crate::serializers::type_serializers::py_err_se_err;
 use crate::tools::{truncate_safe_repr, SchemaDict};
 use crate::PydanticSerializationUnexpectedValue;
@@ -438,10 +437,10 @@ impl TaggedUnionSerializer {
     fn get_discriminator_value(&self, value: &Bound<'_, PyAny>, extra: &Extra) -> Option<Py<PyAny>> {
         let py = value.py();
         let discriminator_value = match &self.discriminator {
-            Discriminator::LookupKey(lookup_key) => match lookup_key {
-                LookupKey::Simple { py_key, .. } => value.getattr(py_key).ok().map(|obj| obj.to_object(py)),
-                _ => None,
-            },
+            Discriminator::LookupKey(lookup_key) => lookup_key
+                .simple_py_get_attr(value)
+                .ok()
+                .and_then(|opt| opt.map(|(_, bound)| bound.to_object(py))),
             Discriminator::Function(func) => func.call1(py, (value,)).ok(),
         };
         if discriminator_value.is_none() {