Add AsNumber for PyUnion

Karatuss · youknowone · commit afb64c873abf · 2023-02-21T19:02:35.000+09:00
diff --git a/vm/src/builtins/union.rs b/vm/src/builtins/union.rs
@@ -6,10 +6,10 @@ use crate::{
     builtins::{PyFrozenSet, PyStr, PyStrRef, PyTuple, PyTupleRef, PyType},
     class::PyClassImpl,
     common::hash,
-    convert::ToPyObject,
+    convert::{ToPyObject, ToPyResult},
     function::PyComparisonValue,
-    protocol::PyMappingMethods,
-    types::{AsMapping, Comparable, GetAttr, Hashable, PyComparisonOp},
+    protocol::{PyMappingMethods, PyNumberMethods},
+    types::{AsMapping, AsNumber, Comparable, GetAttr, Hashable, PyComparisonOp},
     AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, TryFromObject,
     VirtualMachine,
 };
@@ -35,7 +35,7 @@ impl PyPayload for PyUnion {
     }
 }
 
-#[pyclass(with(Hashable, Comparable, AsMapping), flags(BASETYPE))]
+#[pyclass(flags(BASETYPE), with(Hashable, Comparable, AsMapping, AsNumber))]
 impl PyUnion {
     pub fn new(args: PyTupleRef, vm: &VirtualMachine) -> Self {
         let parameters = make_parameters(&args, vm);
@@ -255,6 +255,18 @@ impl AsMapping for PyUnion {
     }
 }
 
+impl AsNumber for PyUnion {
+    fn as_number() -> &'static PyNumberMethods {
+        static AS_NUMBER: Lazy<PyNumberMethods> = Lazy::new(|| PyNumberMethods {
+            or: atomic_func!(|num, other, vm| {
+                PyUnion::or(num.obj.to_owned(), other.to_owned(), vm).to_pyresult(vm)
+            }),
+            ..PyNumberMethods::NOT_IMPLEMENTED
+        });
+        &AS_NUMBER
+    }
+}
+
 impl Comparable for PyUnion {
     fn cmp(
         zelf: &crate::Py<Self>,
