Fix union args (#3864)

Lib/test/test_types.py

@@ -755,8 +755,6 @@ class UnionTests(unittest.TestCase):
         self.assertIs((int | TV)[int], int)
         self.assertIs((TV | int)[int], int)
 
-    # TODO: RUSTPYTHON
-    @unittest.expectedFailure
     def test_union_args(self):
         def check(arg, expected):
             clear_typing_caches()

vm/src/builtins/genericalias.rs

@@ -1,3 +1,4 @@
+use super::type_;
 use crate::{
     builtins::{PyList, PyStr, PyStrRef, PyTuple, PyTupleRef, PyType, PyTypeRef},
     class::PyClassImpl,
@@ -188,6 +189,16 @@ impl PyGenericAlias {
         Err(vm
             .new_type_error("issubclass() argument 2 cannot be a parameterized generic".to_owned()))
     }
+
+    #[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)
+    }
 }
 
 fn is_typevar(obj: &PyObjectRef, vm: &VirtualMachine) -> bool {

vm/src/builtins/type.rs

@@ -403,15 +403,14 @@ impl PyType {
         zelf.iter_mro().map(|cls| cls.clone().into()).collect()
     }
 
-    #[pymethod(name = "__ror__")]
+    #[pymethod(magic)]
+    pub fn ror(zelf: PyObjectRef, other: PyObjectRef, vm: &VirtualMachine) -> PyObjectRef {
+        or_(other, zelf, vm)
+    }
+
     #[pymethod(magic)]
     pub fn or(zelf: PyObjectRef, other: PyObjectRef, vm: &VirtualMachine) -> PyObjectRef {
-        if !union_::is_unionable(zelf.clone(), vm) || !union_::is_unionable(other.clone(), vm) {
-            return vm.ctx.not_implemented();
-        }
-
-        let tuple = PyTuple::new_ref(vec![zelf, other], &vm.ctx);
-        union_::make_union(tuple, vm)
+        or_(zelf, other, vm)
     }
 
     #[pyslot]
@@ -826,6 +825,15 @@ pub(crate) fn call_slot_new(
     unreachable!("Should be able to find a new slot somewhere in the mro")
 }
 
+pub(super) fn or_(zelf: PyObjectRef, other: PyObjectRef, vm: &VirtualMachine) -> PyObjectRef {
+    if !union_::is_unionable(zelf.clone(), vm) || !union_::is_unionable(other.clone(), vm) {
+        return vm.ctx.not_implemented();
+    }
+
+    let tuple = PyTuple::new_ref(vec![zelf, other], &vm.ctx);
+    union_::make_union(tuple, vm)
+}
+
 fn take_next_base(bases: &mut [Vec<PyTypeRef>]) -> Option<PyTypeRef> {
     for base in bases.iter() {
         let head = base[0].clone();

vm/src/builtins/union.rs

@@ -1,4 +1,4 @@
-use super::genericalias;
+use super::{genericalias, type_};
 use crate::{
     builtins::{PyFrozenSet, PyStr, PyStrRef, PyTuple, PyTupleRef, PyType, PyTypeRef},
     class::PyClassImpl,
@@ -100,6 +100,12 @@ impl PyUnion {
         Err(vm
             .new_type_error("issubclass() argument 2 cannot be a parameterized generic".to_owned()))
     }
+
+    #[pymethod(name = "__ror__")]
+    #[pymethod(magic)]
+    fn or(zelf: PyObjectRef, other: PyObjectRef, vm: &VirtualMachine) -> PyObjectRef {
+        type_::or_(zelf, other, vm)
+    }
 }
 
 pub fn is_unionable(obj: PyObjectRef, vm: &VirtualMachine) -> bool {
@@ -156,6 +162,8 @@ fn flatten_args(args: PyTupleRef, vm: &VirtualMachine) -> PyTupleRef {
     for arg in &args {
         if let Some(pyref) = arg.downcast_ref::<PyUnion>() {
             flattened_args.extend(pyref.args.iter().cloned());
+        } else if vm.is_none(arg) {
+            flattened_args.push(vm.ctx.types.none_type.to_owned().into());
         } else {
             flattened_args.push(arg.clone());
         };