diff --git a/vm/src/builtins/object.rs b/vm/src/builtins/object.rs index fa8a88682..c963940a5 100644 --- a/vm/src/builtins/object.rs +++ b/vm/src/builtins/object.rs @@ -174,6 +174,11 @@ impl PyBaseObject { #[pyclassmethod(magic)] fn init_subclass(_cls: PyTypeRef) {} + #[pyclassmethod(magic)] + fn class_getitem(cls: PyTypeRef, _args: FuncArgs) -> PyObjectRef { + cls.into_object() + } + #[pymethod(magic)] pub fn dir(obj: PyObjectRef, vm: &VirtualMachine) -> PyResult { let attributes: PyAttributes = obj.class().get_attributes(); diff --git a/vm/src/pyobject.rs b/vm/src/pyobject.rs index a1598fb5f..5f8ce3c0c 100644 --- a/vm/src/pyobject.rs +++ b/vm/src/pyobject.rs @@ -685,14 +685,26 @@ where T: IntoPyObject, { fn get_item(&self, key: T, vm: &VirtualMachine) -> PyResult { - vm.get_special_method(self.clone(), "__getitem__")? - .map_err(|obj| { - vm.new_type_error(format!( - "'{}' object is not subscriptable", - obj.class().name - )) - })? - .invoke((key,), vm) + match vm.get_special_method(self.clone(), "__getitem__")? { + Ok(special_method) => special_method.invoke((key,), vm), + Err(obj) => { + if obj.isinstance(&vm.ctx.types.type_type) { + vm.get_special_method(obj, "__class_getitem__")? + .map_err(|obj2| { + vm.new_type_error(format!( + "'{}' object is not subscriptable", + obj2.class().name + )) + })? + .invoke((key,), vm) + } else { + Err(vm.new_type_error(format!( + "'{}' object is not subscriptable", + obj.class().name + ))) + } + } + } } fn set_item(&self, key: T, value: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> {