cleanup | implementation

arihant2math · arihant2math · commit f30ec3448d15 · 2025-04-18T21:27:53.000-07:00
Signed-off-by: Ashwin Naren &lt;arihant2math@gmail.com&gt;
diff --git a/vm/src/builtins/genericalias.rs b/vm/src/builtins/genericalias.rs
@@ -1,6 +1,7 @@
 use std::sync::LazyLock;
 
 use super::type_;
+use crate::builtins::union_::UNION_OR;
 use crate::{
     AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, TryFromObject,
     VirtualMachine, atomic_func,
@@ -337,11 +338,7 @@ impl AsMapping for PyGenericAlias {
 
 impl AsNumber for PyGenericAlias {
     fn as_number() -> &'static PyNumberMethods {
-        static AS_NUMBER: PyNumberMethods = PyNumberMethods {
-            or: Some(|a, b, vm| Ok(PyGenericAlias::or(a.to_owned(), b.to_owned(), vm))),
-            ..PyNumberMethods::NOT_IMPLEMENTED
-        };
-        &AS_NUMBER
+        &UNION_OR
     }
 }
 
diff --git a/vm/src/builtins/union.rs b/vm/src/builtins/union.rs
@@ -1,11 +1,12 @@
 use super::{genericalias, type_};
+use crate::class::StaticType;
 use crate::{
     AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine,
     atomic_func,
     builtins::{PyFrozenSet, PyStr, PyTuple, PyTupleRef, PyType},
     class::PyClassImpl,
     common::hash,
-    convert::{ToPyObject, ToPyResult},
+    convert::ToPyObject,
     function::PyComparisonValue,
     protocol::{PyMappingMethods, PyNumberMethods},
     types::{AsMapping, AsNumber, Comparable, GetAttr, Hashable, PyComparisonOp, Representable},
@@ -135,6 +136,9 @@ pub fn is_unionable(obj: PyObjectRef, vm: &VirtualMachine) -> bool {
     obj.class().is(vm.ctx.types.none_type)
         || obj.payload_if_subclass::<PyType>(vm).is_some()
         || obj.class().is(vm.ctx.types.generic_alias_type)
+        || obj
+            .class()
+            .is(crate::stdlib::typing::_typing::TypeAliasType::static_type())
         || obj.class().is(vm.ctx.types.union_type)
 }
 
@@ -229,13 +233,14 @@ impl AsMapping for PyUnion {
     }
 }
 
+pub static UNION_OR: PyNumberMethods = PyNumberMethods {
+    or: Some(|a, b, vm| Ok(type_::or_(a.to_pyobject(vm), b.to_pyobject(vm), vm))),
+    ..PyNumberMethods::NOT_IMPLEMENTED
+};
+
 impl AsNumber for PyUnion {
     fn as_number() -> &'static PyNumberMethods {
-        static AS_NUMBER: PyNumberMethods = PyNumberMethods {
-            or: Some(|a, b, vm| PyUnion::or(a.to_owned(), b.to_owned(), vm).to_pyresult(vm)),
-            ..PyNumberMethods::NOT_IMPLEMENTED
-        };
-        &AS_NUMBER
+        &UNION_OR
     }
 }
 
diff --git a/vm/src/stdlib/typing.rs b/vm/src/stdlib/typing.rs
@@ -2,11 +2,15 @@ pub(crate) use _typing::make_module;
 
 #[pymodule]
 pub(crate) mod _typing {
+    use crate::builtins::union_::UNION_OR;
+    use crate::protocol::PyNumberMethods;
+    use crate::types::AsNumber;
     use crate::{
         PyObjectRef, PyPayload, PyResult, VirtualMachine,
         builtins::{PyGenericAlias, PyTupleRef, PyTypeRef, pystr::AsPyStr},
         function::IntoFuncArgs,
     };
+    use crate::builtins::type_;
 
     pub(crate) fn _call_typing_func_object<'a>(
         _vm: &VirtualMachine,
@@ -37,6 +41,7 @@ pub(crate) mod _typing {
         contravariant: bool,
         infer_variance: bool,
     }
+
     #[pyclass(flags(BASETYPE))]
     impl TypeVar {
         pub(crate) fn _bound(&self, vm: &VirtualMachine) -> PyResult {
@@ -223,7 +228,14 @@ pub(crate) mod _typing {
         // compute_value: PyObjectRef,
         // module: PyObjectRef,
     }
-    #[pyclass(flags(BASETYPE))]
+
+    impl AsNumber for TypeAliasType {
+        fn as_number() -> &'static PyNumberMethods {
+            &UNION_OR
+        }
+    }
+
+    #[pyclass(flags(BASETYPE), with(AsNumber))]
     impl TypeAliasType {
         pub fn new(
             name: PyObjectRef,
@@ -236,6 +248,16 @@ pub(crate) mod _typing {
                 value,
             }
         }
+
+        #[pymethod(magic)]
+        fn ror(zelf: PyObjectRef, other: PyObjectRef, vm: &VirtualMachine) -> PyObjectRef {
+            type_::or_(other, zelf, vm)
+        }
+
+        #[pymethod(magic)]
+        fn or(zelf: PyObjectRef, other: PyObjectRef, vm: &VirtualMachine) -> PyObjectRef {
+            type_::or_(zelf, other, vm)
+        }
     }
 
     #[pyattr]
