Skip to content

Commit 4ebd485

Browse files
authored
typing _Py_subs_paramters (#5880)
1 parent 9f3c34a commit 4ebd485

File tree

2 files changed

+143
-103
lines changed

2 files changed

+143
-103
lines changed

Lib/test/test_typing.py

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -922,8 +922,6 @@ class GenericAliasSubstitutionTests(BaseTestCase):
922922
https://github.com/python/cpython/issues/91162.
923923
"""
924924

925-
# TODO: RUSTPYTHON
926-
@unittest.expectedFailure
927925
def test_one_parameter(self):
928926
T = TypeVar('T')
929927
Ts = TypeVarTuple('Ts')
@@ -1634,8 +1632,6 @@ def test_get_type_hints_on_unpack_args(self):
16341632
# def func3(*args: *CustomVariadic[int, str]): pass
16351633
# self.assertEqual(gth(func3), {'args': Unpack[CustomVariadic[int, str]]})
16361634

1637-
# TODO: RUSTPYTHON
1638-
@unittest.expectedFailure
16391635
def test_get_type_hints_on_unpack_args_string(self):
16401636
Ts = TypeVarTuple('Ts')
16411637

vm/src/builtins/genericalias.rs

Lines changed: 143 additions & 99 deletions
Original file line numberDiff line numberDiff line change
@@ -96,7 +96,22 @@ impl PyGenericAlias {
9696
origin,
9797
args,
9898
parameters,
99-
starred: false, // default to false, will be set to true for Unpack[...]
99+
starred: false, // default to false
100+
}
101+
}
102+
103+
fn with_tuple_args(
104+
origin: PyTypeRef,
105+
args: PyTupleRef,
106+
starred: bool,
107+
vm: &VirtualMachine,
108+
) -> Self {
109+
let parameters = make_parameters(&args, vm);
110+
Self {
111+
origin,
112+
args,
113+
parameters,
114+
starred,
100115
}
101116
}
102117

@@ -166,6 +181,15 @@ impl PyGenericAlias {
166181
self.starred
167182
}
168183

184+
#[pygetset]
185+
fn __typing_unpacked_tuple_args__(&self, vm: &VirtualMachine) -> PyObjectRef {
186+
if self.starred && self.origin.is(vm.ctx.types.tuple_type) {
187+
self.args.clone().into()
188+
} else {
189+
vm.ctx.none()
190+
}
191+
}
192+
169193
#[pymethod]
170194
fn __getitem__(zelf: PyRef<Self>, needle: PyObjectRef, vm: &VirtualMachine) -> PyResult {
171195
let new_args = subs_parameters(
@@ -237,7 +261,7 @@ pub(crate) fn make_parameters(args: &Py<PyTuple>, vm: &VirtualMachine) -> PyTupl
237261
continue;
238262
}
239263

240-
// Check for __typing_subst__ attribute (like CPython)
264+
// Check for __typing_subst__ attribute
241265
if arg.get_attr(identifier!(vm, __typing_subst__), vm).is_ok() {
242266
// Use tuple_add equivalent logic
243267
if tuple_index(&parameters, arg).is_none() {
@@ -336,139 +360,141 @@ fn subs_tvars(
336360
.unwrap_or(Ok(obj))
337361
}
338362

363+
// CPython's _unpack_args equivalent
364+
fn unpack_args(item: PyObjectRef, vm: &VirtualMachine) -> PyResult<PyTupleRef> {
365+
let mut new_args = Vec::new();
366+
367+
let arg_items = if let Ok(tuple) = item.try_to_ref::<PyTuple>(vm) {
368+
tuple.as_slice().to_vec()
369+
} else {
370+
vec![item]
371+
};
372+
373+
for item in arg_items {
374+
// Skip PyType objects - they can't be unpacked
375+
if item.class().is(vm.ctx.types.type_type) {
376+
new_args.push(item);
377+
continue;
378+
}
379+
380+
// Try to get __typing_unpacked_tuple_args__
381+
if let Ok(sub_args) = item.get_attr(identifier!(vm, __typing_unpacked_tuple_args__), vm) {
382+
if !sub_args.is(&vm.ctx.none) {
383+
if let Ok(tuple) = sub_args.try_to_ref::<PyTuple>(vm) {
384+
// Check for ellipsis at the end
385+
let has_ellipsis_at_end = tuple
386+
.as_slice()
387+
.last()
388+
.is_some_and(|item| item.is(&vm.ctx.ellipsis));
389+
390+
if !has_ellipsis_at_end {
391+
// Safe to unpack - add all elements's PyList_SetSlice
392+
for arg in tuple.iter() {
393+
new_args.push(arg.clone());
394+
}
395+
continue;
396+
}
397+
}
398+
}
399+
}
400+
401+
// Default case: add the item as-is's PyList_Append
402+
new_args.push(item);
403+
}
404+
405+
Ok(PyTuple::new_ref(new_args, &vm.ctx))
406+
}
407+
339408
// _Py_subs_parameters
340409
pub fn subs_parameters(
341-
alias: PyObjectRef, // The GenericAlias object itself
410+
alias: PyObjectRef, // = self
342411
args: PyTupleRef,
343412
parameters: PyTupleRef,
344-
needle: PyObjectRef,
413+
item: PyObjectRef,
345414
vm: &VirtualMachine,
346415
) -> PyResult<PyTupleRef> {
347-
let num_params = parameters.len();
348-
if num_params == 0 {
416+
let n_params = parameters.len();
417+
if n_params == 0 {
349418
return Err(vm.new_type_error(format!("{} is not a generic class", alias.repr(vm)?)));
350419
}
351420

352-
// Handle __typing_prepare_subst__ for each parameter
353-
// Following CPython: each prepare function transforms the args
354-
let mut prepared_args = needle.clone();
355-
356-
// Ensure args is a tuple
357-
if prepared_args.try_to_ref::<PyTuple>(vm).is_err() {
358-
prepared_args = PyTuple::new_ref(vec![prepared_args], &vm.ctx).into();
359-
}
421+
// Step 1: Unpack args
422+
let mut item: PyObjectRef = unpack_args(item, vm)?.into();
360423

424+
// Step 2: Call __typing_prepare_subst__ on each parameter
361425
for param in parameters.iter() {
362426
if let Ok(prepare) = param.get_attr(identifier!(vm, __typing_prepare_subst__), vm) {
363427
if !prepare.is(&vm.ctx.none) {
364-
// Call prepare(cls, args) where cls is the GenericAlias
365-
prepared_args = prepare.call((alias.clone(), prepared_args), vm)?;
428+
// Call prepare(self, item)
429+
item = if item.try_to_ref::<PyTuple>(vm).is_ok() {
430+
prepare.call((alias.clone(), item.clone()), vm)?
431+
} else {
432+
// Create a tuple with the single item's "O(O)" format
433+
let tuple_args = PyTuple::new_ref(vec![item.clone()], &vm.ctx);
434+
prepare.call((alias.clone(), tuple_args.to_pyobject(vm)), vm)?
435+
};
366436
}
367437
}
368438
}
369439

370-
let items = prepared_args.try_to_ref::<PyTuple>(vm);
371-
let arg_items = match items {
372-
Ok(tuple) => tuple.as_slice(),
373-
Err(_) => std::slice::from_ref(&prepared_args),
440+
// Step 3: Extract final arg items
441+
let arg_items = if let Ok(tuple) = item.try_to_ref::<PyTuple>(vm) {
442+
tuple.as_slice().to_vec()
443+
} else {
444+
vec![item]
374445
};
446+
let n_items = arg_items.len();
375447

376-
let num_items = arg_items.len();
377-
378-
// Check if we need to apply default values
379-
if num_items < num_params {
380-
// Count how many parameters have defaults
381-
let mut params_with_defaults = 0;
382-
for param in parameters.iter().rev() {
383-
if let Ok(has_default) = vm.call_method(param, "has_default", ()) {
384-
if has_default.try_to_bool(vm)? {
385-
params_with_defaults += 1;
386-
} else {
387-
break; // No more defaults from this point backwards
388-
}
389-
} else {
390-
break;
391-
}
392-
}
393-
394-
let min_required = num_params - params_with_defaults;
395-
if num_items < min_required {
396-
let repr_str = alias.repr(vm)?;
397-
return Err(vm.new_type_error(format!(
398-
"Too few arguments for {repr_str}; actual {num_items}, expected at least {min_required}"
399-
)));
400-
}
401-
} else if num_items > num_params {
402-
let repr_str = alias.repr(vm)?;
448+
if n_items != n_params {
403449
return Err(vm.new_type_error(format!(
404-
"Too many arguments for {repr_str}; actual {num_items}, expected {num_params}"
450+
"Too {} arguments for {}; actual {}, expected {}",
451+
if n_items > n_params { "many" } else { "few" },
452+
alias.repr(vm)?,
453+
n_items,
454+
n_params
405455
)));
406456
}
407457

408-
let mut new_args = Vec::with_capacity(args.len());
458+
// Step 4: Replace all type variables
459+
let mut new_args = Vec::new();
409460

410461
for arg in args.iter() {
411-
// Skip bare Python classes
462+
// Skip PyType objects
412463
if arg.class().is(vm.ctx.types.type_type) {
413464
new_args.push(arg.clone());
414465
continue;
415466
}
416467

417-
// Check if this is an unpacked TypeVarTuple
468+
// Check if this is an unpacked TypeVarTuple's _is_unpacked_typevartuple
418469
let unpack = is_unpacked_typevartuple(arg, vm)?;
419470

420-
// Check for __typing_subst__ attribute directly (like CPython)
421-
if let Ok(subst) = arg.get_attr(identifier!(vm, __typing_subst__), vm) {
422-
if let Some(idx) = tuple_index(parameters.as_slice(), arg) {
423-
if idx < num_items {
424-
// Call __typing_subst__ with the argument
425-
let substituted = subst.call((arg_items[idx].clone(),), vm)?;
426-
427-
if unpack {
428-
// Unpack the tuple if it's a TypeVarTuple
429-
if let Ok(tuple) = substituted.try_to_ref::<PyTuple>(vm) {
430-
for elem in tuple.iter() {
431-
new_args.push(elem.clone());
432-
}
433-
} else {
434-
new_args.push(substituted);
435-
}
436-
} else {
437-
new_args.push(substituted);
438-
}
439-
} else {
440-
// Use default value if available
441-
if let Ok(default_val) = vm.call_method(arg, "__default__", ()) {
442-
if !default_val.is(&vm.ctx.typing_no_default) {
443-
new_args.push(default_val);
444-
} else {
445-
return Err(vm.new_type_error(format!(
446-
"No argument provided for parameter at index {idx}"
447-
)));
448-
}
449-
} else {
450-
return Err(vm.new_type_error(format!(
451-
"No argument provided for parameter at index {idx}"
452-
)));
453-
}
454-
}
471+
// Try __typing_subst__ method first,
472+
let substituted_arg = if let Ok(subst) = arg.get_attr(identifier!(vm, __typing_subst__), vm)
473+
{
474+
// Find parameter index's tuple_index
475+
if let Some(iparam) = tuple_index(parameters.as_slice(), arg) {
476+
subst.call((arg_items[iparam].clone(),), vm)?
455477
} else {
456-
new_args.push(arg.clone());
478+
// This shouldn't happen in well-formed generics but handle gracefully
479+
subs_tvars(arg.clone(), &parameters, &arg_items, vm)?
457480
}
458481
} else {
459-
let subst_arg = subs_tvars(arg.clone(), &parameters, arg_items, vm)?;
460-
if unpack {
461-
// Unpack the tuple if it's a TypeVarTuple
462-
if let Ok(tuple) = subst_arg.try_to_ref::<PyTuple>(vm) {
463-
for elem in tuple.iter() {
464-
new_args.push(elem.clone());
465-
}
466-
} else {
467-
new_args.push(subst_arg);
482+
// Use subs_tvars for objects with __parameters__
483+
subs_tvars(arg.clone(), &parameters, &arg_items, vm)?
484+
};
485+
486+
if unpack {
487+
// Handle unpacked TypeVarTuple's tuple_extend
488+
if let Ok(tuple) = substituted_arg.try_to_ref::<PyTuple>(vm) {
489+
for elem in tuple.iter() {
490+
new_args.push(elem.clone());
468491
}
469492
} else {
470-
new_args.push(subst_arg);
493+
// This shouldn't happen but handle gracefully
494+
new_args.push(substituted_arg);
471495
}
496+
} else {
497+
new_args.push(substituted_arg);
472498
}
473499
}
474500

@@ -565,9 +591,27 @@ impl Representable for PyGenericAlias {
565591
}
566592

567593
impl Iterable for PyGenericAlias {
594+
// ga_iter
595+
// cspell:ignore gaiterobject
596+
// TODO: gaiterobject
568597
fn iter(zelf: PyRef<Self>, vm: &VirtualMachine) -> PyResult {
569-
// Return an iterator over the args tuple
570-
Ok(zelf.args.clone().to_pyobject(vm).get_iter(vm)?.into())
598+
// CPython's ga_iter creates an iterator that yields one starred GenericAlias
599+
// we don't have gaiterobject yet
600+
601+
let starred_alias = PyGenericAlias::with_tuple_args(
602+
zelf.origin.clone(),
603+
zelf.args.clone(),
604+
true, // starred
605+
vm,
606+
);
607+
let starred_ref = PyRef::new_ref(
608+
starred_alias,
609+
vm.ctx.types.generic_alias_type.to_owned(),
610+
None,
611+
);
612+
let items = vec![starred_ref.into()];
613+
let iter_tuple = PyTuple::new_ref(items, &vm.ctx);
614+
Ok(iter_tuple.to_pyobject(vm).get_iter(vm)?.into())
571615
}
572616
}
573617

0 commit comments

Comments
 (0)