Expose `types.GenericAlias.__getitem__`
@@ -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| {
