fix is_enum (pydantic#373)

* fix is_enum as per PyO3/pyo3#2905 * only check is_enum on first loop * logic for pypy 3.7 & 3.8
alexdrydew · Jan 26, 2023 · a478fa4 · a478fa4
1 parent 0ce0f57
commit a478fa4
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 12 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/src/serializers/ob_type.rs b/src/serializers/ob_type.rs
@@ -204,7 +204,7 @@ impl ObTypeLookup {
             ObType::Dataclass
         } else if is_pydantic_model(op_value) {
             ObType::PydanticModel
-        } else if self.is_enum(op_value) {
+        } else if self.is_enum(op_value, type_ptr) {
             ObType::Enum
         } else {
             // this allows for subtypes of the supported class types,
@@ -219,12 +219,17 @@ impl ObTypeLookup {
         }
     }
 
-    fn is_enum(&self, op_value: Option<&PyAny>) -> bool {
-        // waiting for https://github.com/PyO3/pyo3/issues/2905 to be resolved when we should be able to use
-        // `let meta_type = unsafe { (*type_ptr).ob_type };`
-        match op_value {
-            Some(value) => value.get_type().get_type_ptr() as usize == self.enum_type,
-            None => false,
+    fn is_enum(&self, op_value: Option<&PyAny>, type_ptr: *mut PyTypeObject) -> bool {
+        // only test on the type itself, not base types
+        if op_value.is_some() {
+            // see https://github.com/PyO3/pyo3/issues/2905 for details
+            #[cfg(all(PyPy, not(Py_3_9)))]
+            let meta_type = unsafe { (*type_ptr).ob_type };
+            #[cfg(any(not(PyPy), Py_3_9))]
+            let meta_type = unsafe { (*type_ptr).ob_base.ob_base.ob_type };
+            meta_type as usize == self.enum_type
+        } else {
+            false
         }
     }
 }