Repository: RustPython/RustPython
File: crates/vm/src/builtins/union.rs (Rust)

Code before:
    }
    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)
        && (!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)?;
    Ok(result.args)
}

fn flatten_args(args: &Py<PyTuple>, vm: &VirtualMachine) -> PyTupleRef {
    let mut total_args = 0;
    for arg in args {
        if let Some(pyref) = arg.downcast_ref::<PyUnion>() {
            total_args += pyref.args.len();
        } else {
            total_args += 1;

Review comment (suggestion):
```suggestion
    let left = type_check(zelf, vm)?;
```

Code after:
}

pub fn or_op(zelf: PyObjectRef, other: PyObjectRef, vm: &VirtualMachine) -> PyResult {
    if !has_union_operands(zelf.clone(), other.clone(), vm)
        && (!is_unionable(zelf.clone(), vm) || !is_unionable(other.clone(), vm))
    {
        return Ok(vm.ctx.not_implemented());
    }

    let left = type_check(zelf, vm)?;
    let right = type_check(other, vm)?;
    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)?;
    Ok(result.args)
}

fn flatten_args(args: &Py<PyTuple>, vm: &VirtualMachine) -> PyTupleRef {
    let mut total_args = 0;
    for arg in args {
        if let Some(pyref) = arg.downcast_ref::<PyUnion>() {
            total_args += pyref.args.len();
        } else {
            total_args += 1;
        };
    }

    let mut flattened_args = Vec::with_capacity(total_args);
    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 if arg.downcast_ref::<PyStr>().is_some() {
            // Convert string to ForwardRef
            match string_to_forwardref(arg.clone(), vm) {
                Ok(fr) => flattened_args.push(fr),
                Err(_) => flattened_args.push(arg.clone()),
            }
        } else {
            flattened_args.push(arg.clone());
        };
    }

    PyTuple::new_ref(flattened_args, &vm.ctx)
}