From 523a0f9bc688d1dad873889d685ca7fced56099a Mon Sep 17 00:00:00 2001
From: Karatus <shm1193@gmail.com>
Date: Wed, 13 Jul 2022 18:45:02 +0900
Subject: [PATCH] Fix union args (#3864)

---
 Lib/test/test_types.py          |  2 --
 vm/src/builtins/genericalias.rs | 11 +++++++++++
 vm/src/builtins/type.rs         | 22 +++++++++++++++-------
 vm/src/builtins/union.rs        | 10 +++++++++-
 4 files changed, 35 insertions(+), 10 deletions(-)

diff --git a/Lib/test/test_types.py b/Lib/test/test_types.py
index d314387d6..9ba6a077e 100644
--- a/Lib/test/test_types.py
+++ b/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()
diff --git a/vm/src/builtins/genericalias.rs b/vm/src/builtins/genericalias.rs
index 4f1d6fedd..7e0c061e7 100644
--- a/vm/src/builtins/genericalias.rs
+++ b/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 {
diff --git a/vm/src/builtins/type.rs b/vm/src/builtins/type.rs
index d4e09f8c2..12de79545 100644
--- a/vm/src/builtins/type.rs
+++ b/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();
diff --git a/vm/src/builtins/union.rs b/vm/src/builtins/union.rs
index 58d9ae4dc..d43fbc10b 100644
--- a/vm/src/builtins/union.rs
+++ b/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());
         };