diff --git a/vm/src/builtins/dict.rs b/vm/src/builtins/dict.rs index 1803fdf17..514dd54f5 100644 --- a/vm/src/builtins/dict.rs +++ b/vm/src/builtins/dict.rs @@ -1092,7 +1092,15 @@ trait ViewSetOps: DictView { } impl ViewSetOps for PyDictKeys {} -#[pyclass(with(DictView, Constructor, Comparable, Iterable, ViewSetOps, AsSequence))] +#[pyclass(with( + DictView, + Constructor, + Comparable, + Iterable, + ViewSetOps, + AsSequence, + AsNumber +))] impl PyDictKeys { #[pymethod(magic)] fn contains(zelf: PyRef, key: PyObjectRef, vm: &VirtualMachine) -> PyResult { @@ -1133,8 +1141,70 @@ impl AsSequence for PyDictKeys { } } +impl AsNumber for PyDictKeys { + fn as_number() -> &'static PyNumberMethods { + static AS_NUMBER: PyNumberMethods = PyNumberMethods { + subtract: Some(|num, args, vm| { + let num = PySetInner::from_iter( + ArgIterable::try_from_object(vm, num.obj.to_owned())?.iter(vm)?, + vm, + )?; + Ok(PySet { + inner: num + .difference(ArgIterable::try_from_object(vm, args.to_owned())?, vm)?, + } + .into_pyobject(vm)) + }), + and: Some(|num, args, vm| { + let num = PySetInner::from_iter( + ArgIterable::try_from_object(vm, num.obj.to_owned())?.iter(vm)?, + vm, + )?; + Ok(PySet { + inner: num + .intersection(ArgIterable::try_from_object(vm, args.to_owned())?, vm)?, + } + .into_pyobject(vm)) + }), + xor: Some(|num, args, vm| { + let num = PySetInner::from_iter( + ArgIterable::try_from_object(vm, num.obj.to_owned())?.iter(vm)?, + vm, + )?; + Ok(PySet { + inner: num.symmetric_difference( + ArgIterable::try_from_object(vm, args.to_owned())?, + vm, + )?, + } + .into_pyobject(vm)) + }), + or: Some(|num, args, vm| { + let num = PySetInner::from_iter( + ArgIterable::try_from_object(vm, num.obj.to_owned())?.iter(vm)?, + vm, + )?; + Ok(PySet { + inner: num.union(ArgIterable::try_from_object(vm, args.to_owned())?, vm)?, + } + .into_pyobject(vm)) + }), + ..PyNumberMethods::NOT_IMPLEMENTED + }; + &AS_NUMBER + } +} + impl ViewSetOps for PyDictItems {} -#[pyclass(with(DictView, Constructor, Comparable, Iterable, ViewSetOps, AsSequence))] +#[pyclass(with( + DictView, + Constructor, + Comparable, + Iterable, + ViewSetOps, + AsSequence, + AsNumber +))] impl PyDictItems { #[pymethod(magic)] fn contains(zelf: PyRef, needle: PyObjectRef, vm: &VirtualMachine) -> PyResult { @@ -1189,6 +1259,60 @@ impl AsSequence for PyDictItems { } } +impl AsNumber for PyDictItems { + fn as_number() -> &'static PyNumberMethods { + static AS_NUMBER: PyNumberMethods = PyNumberMethods { + subtract: Some(|num, args, vm| { + let num = PySetInner::from_iter( + ArgIterable::try_from_object(vm, num.obj.to_owned())?.iter(vm)?, + vm, + )?; + Ok(PySet { + inner: num + .difference(ArgIterable::try_from_object(vm, args.to_owned())?, vm)?, + } + .into_pyobject(vm)) + }), + and: Some(|num, args, vm| { + let num = PySetInner::from_iter( + ArgIterable::try_from_object(vm, num.obj.to_owned())?.iter(vm)?, + vm, + )?; + Ok(PySet { + inner: num + .intersection(ArgIterable::try_from_object(vm, args.to_owned())?, vm)?, + } + .into_pyobject(vm)) + }), + xor: Some(|num, args, vm| { + let num = PySetInner::from_iter( + ArgIterable::try_from_object(vm, num.obj.to_owned())?.iter(vm)?, + vm, + )?; + Ok(PySet { + inner: num.symmetric_difference( + ArgIterable::try_from_object(vm, args.to_owned())?, + vm, + )?, + } + .into_pyobject(vm)) + }), + or: Some(|num, args, vm| { + let num = PySetInner::from_iter( + ArgIterable::try_from_object(vm, num.obj.to_owned())?.iter(vm)?, + vm, + )?; + Ok(PySet { + inner: num.union(ArgIterable::try_from_object(vm, args.to_owned())?, vm)?, + } + .into_pyobject(vm)) + }), + ..PyNumberMethods::NOT_IMPLEMENTED + }; + &AS_NUMBER + } +} + #[pyclass(with(DictView, Constructor, Iterable, AsSequence))] impl PyDictValues { #[pygetset]