support | operation between typing.Union and strings

elmjag · elmjag · commit 7d0eeddfcfe6 · 2026-02-03T14:31:06.000+01:00
Adds support for performing '|' operation between Union objects and
strings, e.g. forward type references.

For example following code:

    from typing import Union

    U1 = Union[int, str]
    U1 | "float"

The result of the operation above becomes:

   int | str | ForwardRef('float')
diff --git a/Lib/test/test_typing.py b/Lib/test/test_typing.py
@@ -2281,7 +2281,6 @@ class Ints(enum.IntEnum):
         self.assertEqual(Union[Literal[1], Literal[Ints.B], Literal[True]].__args__,
                          (Literal[1], Literal[Ints.B], Literal[True]))
 
-    @unittest.expectedFailure  # TODO: RUSTPYTHON; AssertionError: types.UnionType[int, str] | float != types.UnionType[int, str, float]
     def test_allow_non_types_in_or(self):
         # gh-140348: Test that using | with a Union object allows things that are
         # not allowed by is_unionable().
diff --git a/crates/vm/src/builtins/type.rs b/crates/vm/src/builtins/type.rs
@@ -2048,12 +2048,7 @@ pub(crate) fn call_slot_new(
 }
 
 pub(crate) fn or_(zelf: PyObjectRef, other: PyObjectRef, vm: &VirtualMachine) -> PyResult {
-    if !union_::is_unionable(zelf.clone(), vm) || !union_::is_unionable(other.clone(), vm) {
-        return Ok(vm.ctx.not_implemented());
-    }
-
-    let tuple = PyTuple::new_ref(vec![zelf, other], &vm.ctx);
-    union_::make_union(&tuple, vm)
+    union_::or_op(zelf, other, vm)
 }
 
 fn take_next_base(bases: &mut [Vec<PyTypeRef>]) -> Option<PyTypeRef> {
diff --git a/crates/vm/src/builtins/union.rs b/crates/vm/src/builtins/union.rs
@@ -2,10 +2,11 @@ use super::{genericalias, type_};
 use crate::{
     AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine,
     atomic_func,
-    builtins::{PyFrozenSet, PySet, PyStr, PyTuple, PyTupleRef, PyType},
+    builtins::{PyFrozenSet, PySet, PyStr, PyTuple, PyTupleRef, PyType, pystr::AsPyStr},
     class::PyClassImpl,
     common::hash,
     convert::ToPyObject,
+    function::IntoFuncArgs,
     function::PyComparisonValue,
     protocol::{PyMappingMethods, PyNumberMethods},
     stdlib::typing::TypeAliasType,
@@ -193,7 +194,7 @@ impl PyUnion {
     }
 }
 
-pub fn is_unionable(obj: PyObjectRef, vm: &VirtualMachine) -> bool {
+fn is_unionable(obj: PyObjectRef, vm: &VirtualMachine) -> bool {
     let cls = obj.class();
     cls.is(vm.ctx.types.none_type)
         || obj.downcastable::<PyType>()
@@ -202,6 +203,58 @@ pub fn is_unionable(obj: PyObjectRef, vm: &VirtualMachine) -> bool {
         || obj.downcast_ref::<TypeAliasType>().is_some()
 }
 
+fn _call_typing_func_object<'a>(
+    vm: &VirtualMachine,
+    func_name: impl AsPyStr<'a>,
+    args: impl IntoFuncArgs,
+) -> PyResult {
+    let module = vm.import("typing", 0)?;
+    let func = module.get_attr(func_name.as_pystr(&vm.ctx), vm)?;
+    func.call(args, vm)
+}
+
+fn type_check(arg: PyObjectRef, vm: &VirtualMachine) -> PyResult<PyObjectRef> {
+    // Fast path to avoid calling into typing.py
+    if is_unionable(arg.clone(), vm) {
+        return Ok(arg);
+    }
+    let message_str: PyObjectRef = vm
+        .ctx
+        .new_str("Union[arg, ...]: each arg must be a type.")
+        .into();
+    _call_typing_func_object(vm, "_type_check", (arg, message_str))
+}
+
+fn has_union_operands(a: PyObjectRef, b: PyObjectRef, vm: &VirtualMachine) -> bool {
+    let union_type = vm.ctx.types.union_type;
+    a.class().is(union_type) || b.class().is(union_type)
+}
+
+pub fn or_op(zelf: PyObjectRef, other: PyObjectRef, vm: &VirtualMachine) -> PyResult {
+    if !has_union_operands(zelf.clone(), other.clone(), vm) {
+        if !is_unionable(zelf.clone(), vm) || !is_unionable(other.clone(), vm) {
+            return Ok(vm.ctx.not_implemented());
+        }
+    }
+
+    let left = match type_check(zelf, vm) {
+        Ok(v) => v,
+        err => {
+            return err;
+        }
+    };
+
+    let right = match type_check(other, vm) {
+        Ok(v) => v,
+        err => {
+            return err;
+        }
+    };
+
+    let tuple = PyTuple::new_ref(vec![left, right], &vm.ctx);
+    make_union(&tuple, vm)
+}
+
 fn make_parameters(args: &Py<PyTuple>, vm: &VirtualMachine) -> PyResult<PyTupleRef> {
     let parameters = genericalias::make_parameters(args, vm);
     let result = dedup_and_flatten_args(&parameters, vm)?;

Original file line number	Diff line number	Diff line change
`@@ -2048,12 +2048,7 @@ pub(crate) fn call_slot_new(`
`2048`	`2048`	`}`
`2049`	`2049`
`2050`	`2050`	`pub(crate) fn or_(zelf: PyObjectRef, other: PyObjectRef, vm: &VirtualMachine) -> PyResult {`
`2051`		`- if !union_::is_unionable(zelf.clone(), vm) \|\| !union_::is_unionable(other.clone(), vm) {`
`2052`		`- return Ok(vm.ctx.not_implemented());`
`2053`		`- }`
`2054`		`-`
`2055`		`- let tuple = PyTuple::new_ref(vec![zelf, other], &vm.ctx);`
`2056`		`- union_::make_union(&tuple, vm)`
	`2051`	`+ union_::or_op(zelf, other, vm)`
`2057`	`2052`	`}`
`2058`	`2053`
`2059`	`2054`	`fn take_next_base(bases: &mut [Vec<PyTypeRef>]) -> Option<PyTypeRef> {`