Skip to content

Commit 8dbc58c

Browse files
committed
typing _Py_subs_paramters
1 parent 9f3c34a commit 8dbc58c

File tree

2 files changed

+144
-103
lines changed

2 files changed

+144
-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: 144 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,142 @@ 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 = if tuple.len() > 0 {
386+
tuple.as_slice()[tuple.len() - 1].is(&vm.ctx.ellipsis)
387+
} else {
388+
false
389+
};
390+
391+
if !has_ellipsis_at_end {
392+
// Safe to unpack - add all elements's PyList_SetSlice
393+
for arg in tuple.iter() {
394+
new_args.push(arg.clone());
395+
}
396+
continue;
397+
}
398+
}
399+
}
400+
}
401+
402+
// Default case: add the item as-is's PyList_Append
403+
new_args.push(item);
404+
}
405+
406+
Ok(PyTuple::new_ref(new_args, &vm.ctx))
407+
}
408+
339409
// _Py_subs_parameters
340410
pub fn subs_parameters(
341-
alias: PyObjectRef, // The GenericAlias object itself
411+
alias: PyObjectRef, // = self
342412
args: PyTupleRef,
343413
parameters: PyTupleRef,
344-
needle: PyObjectRef,
414+
item: PyObjectRef,
345415
vm: &VirtualMachine,
346416
) -> PyResult<PyTupleRef> {
347-
let num_params = parameters.len();
348-
if num_params == 0 {
417+
let n_params = parameters.len();
418+
if n_params == 0 {
349419
return Err(vm.new_type_error(format!("{} is not a generic class", alias.repr(vm)?)));
350420
}
351421

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-
}
422+
// Step 1: Unpack args
423+
let mut item: PyObjectRef = unpack_args(item, vm)?.into();
360424

425+
// Step 2: Call __typing_prepare_subst__ on each parameter
361426
for param in parameters.iter() {
362427
if let Ok(prepare) = param.get_attr(identifier!(vm, __typing_prepare_subst__), vm) {
363428
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)?;
429+
// Call prepare(self, item)
430+
item = if item.try_to_ref::<PyTuple>(vm).is_ok() {
431+
prepare.call((alias.clone(), item.clone()), vm)?
432+
} else {
433+
// Create a tuple with the single item's "O(O)" format
434+
let tuple_args = PyTuple::new_ref(vec![item.clone()], &vm.ctx);
435+
prepare.call((alias.clone(), tuple_args.to_pyobject(vm)), vm)?
436+
};
366437
}
367438
}
368439
}
369440

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),
441+
// Step 3: Extract final arg items
442+
let arg_items = if let Ok(tuple) = item.try_to_ref::<PyTuple>(vm) {
443+
tuple.as_slice().to_vec()
444+
} else {
445+
vec![item]
374446
};
447+
let n_items = arg_items.len();
375448

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)?;
449+
if n_items != n_params {
403450
return Err(vm.new_type_error(format!(
404-
"Too many arguments for {repr_str}; actual {num_items}, expected {num_params}"
451+
"Too {} arguments for {}; actual {}, expected {}",
452+
if n_items > n_params { "many" } else { "few" },
453+
alias.repr(vm)?,
454+
n_items,
455+
n_params
405456
)));
406457
}
407458

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

410462
for arg in args.iter() {
411-
// Skip bare Python classes
463+
// Skip PyType objects
412464
if arg.class().is(vm.ctx.types.type_type) {
413465
new_args.push(arg.clone());
414466
continue;
415467
}
416468

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

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-
}
472+
// Try __typing_subst__ method first,
473+
let substituted_arg = if let Ok(subst) = arg.get_attr(identifier!(vm, __typing_subst__), vm)
474+
{
475+
// Find parameter index's tuple_index
476+
if let Some(iparam) = tuple_index(parameters.as_slice(), arg) {
477+
subst.call((arg_items[iparam].clone(),), vm)?
455478
} else {
456-
new_args.push(arg.clone());
479+
// This shouldn't happen in well-formed generics but handle gracefully
480+
subs_tvars(arg.clone(), &parameters, &arg_items, vm)?
457481
}
458482
} 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);
483+
// Use subs_tvars for objects with __parameters__
484+
subs_tvars(arg.clone(), &parameters, &arg_items, vm)?
485+
};
486+
487+
if unpack {
488+
// Handle unpacked TypeVarTuple's tuple_extend
489+
if let Ok(tuple) = substituted_arg.try_to_ref::<PyTuple>(vm) {
490+
for elem in tuple.iter() {
491+
new_args.push(elem.clone());
468492
}
469493
} else {
470-
new_args.push(subst_arg);
494+
// This shouldn't happen but handle gracefully
495+
new_args.push(substituted_arg);
471496
}
497+
} else {
498+
new_args.push(substituted_arg);
472499
}
473500
}
474501

@@ -565,9 +592,27 @@ impl Representable for PyGenericAlias {
565592
}
566593

567594
impl Iterable for PyGenericAlias {
595+
// ga_iter
596+
// cspell:ignore gaiterobject
597+
// TODO: gaiterobject
568598
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())
599+
// CPython's ga_iter creates an iterator that yields one starred GenericAlias
600+
// we don't have gaiterobject yet
601+
602+
let starred_alias = PyGenericAlias::with_tuple_args(
603+
zelf.origin.clone(),
604+
zelf.args.clone(),
605+
true, // starred
606+
vm,
607+
);
608+
let starred_ref = PyRef::new_ref(
609+
starred_alias,
610+
vm.ctx.types.generic_alias_type.to_owned(),
611+
None,
612+
);
613+
let items = vec![starred_ref.into()];
614+
let iter_tuple = PyTuple::new_ref(items, &vm.ctx);
615+
Ok(iter_tuple.to_pyobject(vm).get_iter(vm)?.into())
571616
}
572617
}
573618

0 commit comments

Comments
 (0)