Skip to content

Commit ffc4622

Browse files
elmjagelmjag
authored
support | operation between typing.Union and strings (#6983)
* remove duplicated _call_typing_func_object() functions Move _call_typing_func_object() code to stdlib::typing::call_typing_func_object(). Use that function everywhere. * support | operation between typing.Union and strings 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')
1 parent a037bda commit ffc4622

File tree

5 files changed

+57
-39
lines changed

5 files changed

+57
-39
lines changed

‎Lib/test/test_typing.py‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2281,7 +2281,6 @@ class Ints(enum.IntEnum):
22812281
self.assertEqual(Union[Literal[1], Literal[Ints.B], Literal[True]].__args__,
22822282
(Literal[1], Literal[Ints.B], Literal[True]))
22832283

2284-
@unittest.expectedFailure # TODO: RUSTPYTHON; AssertionError: types.UnionType[int, str] | float != types.UnionType[int, str, float]
22852284
def test_allow_non_types_in_or(self):
22862285
# gh-140348: Test that using | with a Union object allows things that are
22872286
# not allowed by is_unionable().

‎crates/vm/src/builtins/type.rs‎

Lines changed: 1 addition & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -2048,12 +2048,7 @@ pub(crate) fn call_slot_new(
20482048
}
20492049

20502050
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)
20572052
}
20582053

20592054
fn take_next_base(bases: &mut [Vec<PyTypeRef>]) -> Option<PyTypeRef> {

‎crates/vm/src/builtins/union.rs‎

Lines changed: 32 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ use crate::{
88
convert::ToPyObject,
99
function::PyComparisonValue,
1010
protocol::{PyMappingMethods, PyNumberMethods},
11-
stdlib::typing::TypeAliasType,
11+
stdlib::typing::{TypeAliasType, call_typing_func_object},
1212
types::{AsMapping, AsNumber, Comparable, GetAttr, Hashable, PyComparisonOp, Representable},
1313
};
1414
use alloc::fmt;
@@ -193,7 +193,7 @@ impl PyUnion {
193193
}
194194
}
195195

196-
pub fn is_unionable(obj: PyObjectRef, vm: &VirtualMachine) -> bool {
196+
fn is_unionable(obj: PyObjectRef, vm: &VirtualMachine) -> bool {
197197
let cls = obj.class();
198198
cls.is(vm.ctx.types.none_type)
199199
|| obj.downcastable::<PyType>()
@@ -202,6 +202,36 @@ pub fn is_unionable(obj: PyObjectRef, vm: &VirtualMachine) -> bool {
202202
|| obj.downcast_ref::<TypeAliasType>().is_some()
203203
}
204204

205+
fn type_check(arg: PyObjectRef, vm: &VirtualMachine) -> PyResult<PyObjectRef> {
206+
// Fast path to avoid calling into typing.py
207+
if is_unionable(arg.clone(), vm) {
208+
return Ok(arg);
209+
}
210+
let message_str: PyObjectRef = vm
211+
.ctx
212+
.new_str("Union[arg, ...]: each arg must be a type.")
213+
.into();
214+
call_typing_func_object(vm, "_type_check", (arg, message_str))
215+
}
216+
217+
fn has_union_operands(a: PyObjectRef, b: PyObjectRef, vm: &VirtualMachine) -> bool {
218+
let union_type = vm.ctx.types.union_type;
219+
a.class().is(union_type) || b.class().is(union_type)
220+
}
221+
222+
pub fn or_op(zelf: PyObjectRef, other: PyObjectRef, vm: &VirtualMachine) -> PyResult {
223+
if !has_union_operands(zelf.clone(), other.clone(), vm)
224+
&& (!is_unionable(zelf.clone(), vm) || !is_unionable(other.clone(), vm))
225+
{
226+
return Ok(vm.ctx.not_implemented());
227+
}
228+
229+
let left = type_check(zelf, vm)?;
230+
let right = type_check(other, vm)?;
231+
let tuple = PyTuple::new_ref(vec![left, right], &vm.ctx);
232+
make_union(&tuple, vm)
233+
}
234+
205235
fn make_parameters(args: &Py<PyTuple>, vm: &VirtualMachine) -> PyResult<PyTupleRef> {
206236
let parameters = genericalias::make_parameters(args, vm);
207237
let result = dedup_and_flatten_args(&parameters, vm)?;

‎crates/vm/src/stdlib/typevar.rs‎

Lines changed: 8 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -6,30 +6,21 @@ pub use typevar::*;
66
pub(crate) mod typevar {
77
use crate::{
88
AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine,
9-
builtins::{PyTuple, PyTupleRef, PyType, PyTypeRef, make_union, pystr::AsPyStr},
9+
builtins::{PyTuple, PyTupleRef, PyType, PyTypeRef, make_union},
1010
common::lock::PyMutex,
11-
function::{FuncArgs, IntoFuncArgs, PyComparisonValue},
11+
function::{FuncArgs, PyComparisonValue},
1212
protocol::PyNumberMethods,
13+
stdlib::typing::call_typing_func_object,
1314
types::{AsNumber, Comparable, Constructor, Iterable, PyComparisonOp, Representable},
1415
};
1516

16-
pub(crate) fn _call_typing_func_object<'a>(
17-
vm: &VirtualMachine,
18-
func_name: impl AsPyStr<'a>,
19-
args: impl IntoFuncArgs,
20-
) -> PyResult {
21-
let module = vm.import("typing", 0)?;
22-
let func = module.get_attr(func_name.as_pystr(&vm.ctx), vm)?;
23-
func.call(args, vm)
24-
}
25-
2617
fn type_check(arg: PyObjectRef, msg: &str, vm: &VirtualMachine) -> PyResult<PyObjectRef> {
2718
// Calling typing.py here leads to bootstrapping problems
2819
if vm.is_none(&arg) {
2920
return Ok(arg.class().to_owned().into());
3021
}
3122
let message_str: PyObjectRef = vm.ctx.new_str(msg).into();
32-
_call_typing_func_object(vm, "_type_check", (arg, message_str))
23+
call_typing_func_object(vm, "_type_check", (arg, message_str))
3324
}
3425

3526
/// Get the module of the caller frame, similar to CPython's caller() function.
@@ -169,7 +160,7 @@ pub(crate) mod typevar {
169160
vm: &VirtualMachine,
170161
) -> PyResult {
171162
let self_obj: PyObjectRef = zelf.into();
172-
_call_typing_func_object(vm, "_typevar_subst", (self_obj, arg))
163+
call_typing_func_object(vm, "_typevar_subst", (self_obj, arg))
173164
}
174165

175166
#[pymethod]
@@ -514,7 +505,7 @@ pub(crate) mod typevar {
514505
vm: &VirtualMachine,
515506
) -> PyResult {
516507
let self_obj: PyObjectRef = zelf.into();
517-
_call_typing_func_object(vm, "_paramspec_subst", (self_obj, arg))
508+
call_typing_func_object(vm, "_paramspec_subst", (self_obj, arg))
518509
}
519510

520511
#[pymethod]
@@ -525,7 +516,7 @@ pub(crate) mod typevar {
525516
vm: &VirtualMachine,
526517
) -> PyResult {
527518
let self_obj: PyObjectRef = zelf.into();
528-
_call_typing_func_object(vm, "_paramspec_prepare_subst", (self_obj, alias, args))
519+
call_typing_func_object(vm, "_paramspec_prepare_subst", (self_obj, alias, args))
529520
}
530521
}
531522

@@ -711,7 +702,7 @@ pub(crate) mod typevar {
711702
vm: &VirtualMachine,
712703
) -> PyResult {
713704
let self_obj: PyObjectRef = zelf.into();
714-
_call_typing_func_object(vm, "_typevartuple_prepare_subst", (self_obj, alias, args))
705+
call_typing_func_object(vm, "_typevartuple_prepare_subst", (self_obj, alias, args))
715706
}
716707
}
717708

‎crates/vm/src/stdlib/typing.rs‎

Lines changed: 16 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,8 @@
11
// spell-checker:ignore typevarobject funcobj
2-
use crate::{Context, class::PyClassImpl};
2+
use crate::{
3+
Context, PyResult, VirtualMachine, builtins::pystr::AsPyStr, class::PyClassImpl,
4+
function::IntoFuncArgs,
5+
};
36

47
pub use crate::stdlib::typevar::{
58
Generic, ParamSpec, ParamSpecArgs, ParamSpecKwargs, TypeVar, TypeVarTuple,
@@ -13,26 +16,26 @@ pub fn init(ctx: &Context) {
1316
NoDefault::extend_class(ctx, ctx.types.typing_no_default_type);
1417
}
1518

19+
pub fn call_typing_func_object<'a>(
20+
vm: &VirtualMachine,
21+
func_name: impl AsPyStr<'a>,
22+
args: impl IntoFuncArgs,
23+
) -> PyResult {
24+
let module = vm.import("typing", 0)?;
25+
let func = module.get_attr(func_name.as_pystr(&vm.ctx), vm)?;
26+
func.call(args, vm)
27+
}
28+
1629
#[pymodule(name = "_typing", with(super::typevar::typevar))]
1730
pub(crate) mod decl {
1831
use crate::{
1932
Py, PyObjectRef, PyPayload, PyResult, VirtualMachine,
20-
builtins::{PyStrRef, PyTupleRef, PyType, PyTypeRef, pystr::AsPyStr, type_},
21-
function::{FuncArgs, IntoFuncArgs},
33+
builtins::{PyStrRef, PyTupleRef, PyType, PyTypeRef, type_},
34+
function::FuncArgs,
2235
protocol::PyNumberMethods,
2336
types::{AsNumber, Constructor, Representable},
2437
};
2538

26-
pub(crate) fn _call_typing_func_object<'a>(
27-
vm: &VirtualMachine,
28-
func_name: impl AsPyStr<'a>,
29-
args: impl IntoFuncArgs,
30-
) -> PyResult {
31-
let module = vm.import("typing", 0)?;
32-
let func = module.get_attr(func_name.as_pystr(&vm.ctx), vm)?;
33-
func.call(args, vm)
34-
}
35-
3639
#[pyfunction]
3740
pub(crate) fn _idfunc(args: FuncArgs, _vm: &VirtualMachine) -> PyObjectRef {
3841
args.args[0].clone()

0 commit comments

Comments
 (0)