Merge pull request #3499 from moreal/bugfix/genericalias-getitem-method

Expose `types.GenericAlias.__getitem__`
This commit is contained in:
Jeong YunWon
2022-01-22 15:50:21 +09:00
committed by GitHub

View File

@@ -126,6 +126,58 @@ impl PyGenericAlias {
self.origin.as_object().to_owned()
}
#[pymethod(magic)]
fn getitem(&self, needle: PyObjectRef, vm: &VirtualMachine) -> PyResult {
let num_params = self.parameters.len();
if num_params == 0 {
return Err(vm.new_type_error(format!(
"There are no type variables left in {}",
self.repr(vm)?
)));
}
let items = PyTupleRef::try_from_object(vm, needle.clone());
let arg_items = match items {
Ok(ref tuple) => tuple.as_slice(),
Err(_) => std::slice::from_ref(&needle),
};
let num_items = arg_items.len();
if num_params != num_items {
let plural = if num_items > num_params {
"many"
} else {
"few"
};
return Err(vm.new_type_error(format!(
"Too {} arguments for {}",
plural,
self.repr(vm)?
)));
}
let new_args = self
.args
.as_slice()
.iter()
.map(|arg| {
if is_typevar(arg) {
let idx = tuple_index(&self.parameters, arg).unwrap();
Ok(arg_items[idx].clone())
} else {
subs_tvars(arg.clone(), &self.parameters, arg_items, vm)
}
})
.collect::<PyResult<Vec<_>>>()?;
Ok(PyGenericAlias::new(
self.origin.clone(),
PyTuple::new_ref(new_args, &vm.ctx).into(),
vm,
)
.into_object(vm))
}
#[pymethod(magic)]
fn dir(&self, vm: &VirtualMachine) -> PyResult<PyList> {
let dir = vm.dir(Some(self.origin()))?;
@@ -238,57 +290,6 @@ fn subs_tvars(
}
impl PyGenericAlias {
fn getitem(&self, needle: PyObjectRef, vm: &VirtualMachine) -> PyResult {
let num_params = self.parameters.len();
if num_params == 0 {
return Err(vm.new_type_error(format!(
"There are no type variables left in {}",
self.repr(vm)?
)));
}
let items = PyTupleRef::try_from_object(vm, needle.clone());
let arg_items = match items {
Ok(ref tuple) => tuple.as_slice(),
Err(_) => std::slice::from_ref(&needle),
};
let num_items = arg_items.len();
if num_params != num_items {
let plural = if num_items > num_params {
"many"
} else {
"few"
};
return Err(vm.new_type_error(format!(
"Too {} arguments for {}",
plural,
self.repr(vm)?
)));
}
let new_args = self
.args
.as_slice()
.iter()
.map(|arg| {
if is_typevar(arg) {
let idx = tuple_index(&self.parameters, arg).unwrap();
Ok(arg_items[idx].clone())
} else {
subs_tvars(arg.clone(), &self.parameters, arg_items, vm)
}
})
.collect::<PyResult<Vec<_>>>()?;
Ok(PyGenericAlias::new(
self.origin.clone(),
PyTuple::new_ref(new_args, &vm.ctx).into(),
vm,
)
.into_object(vm))
}
const MAPPING_METHODS: PyMappingMethods = PyMappingMethods {
length: None,
subscript: Some(|mapping, needle, vm| {