Skip to content

Commit 2d02933

Browse files
youknowoneyouknowone
authored
Merge pull request #4619 from xiaozhiyan/implement-number-protocol-for-pytype-pyset-pyfrozenset
Implement Number Protocol for `PyType`, `PySet` and `PyFrozenSet`
2 parents c656cdd + 1bec1d5 commit 2d02933

File tree

2 files changed

+128
-14
lines changed

2 files changed

+128
-14
lines changed

‎vm/src/builtins/set.rs

+106-7
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,20 @@
1-
use once_cell::sync::Lazy;
2-
31
/*
42
* Builtin set type with a sequence of unique items.
53
*/
64
use super::{
75
builtins_iter, IterStatus, PositionIterInternal, PyDict, PyDictRef, PyGenericAlias, PyTupleRef,
86
PyType, PyTypeRef,
97
};
10-
use crate::atomic_func;
11-
use crate::common::{ascii, hash::PyHash, lock::PyMutex, rc::PyRc};
128
use crate::{
9+
atomic_func,
1310
class::PyClassImpl,
11+
common::{ascii, hash::PyHash, lock::PyMutex, rc::PyRc},
12+
convert::ToPyResult,
1413
dictdatatype::{self, DictSize},
1514
function::{ArgIterable, FuncArgs, OptionalArg, PosArgs, PyArithmeticValue, PyComparisonValue},
16-
protocol::{PyIterReturn, PySequenceMethods},
15+
protocol::{PyIterReturn, PyNumberMethods, PySequenceMethods},
1716
recursion::ReprGuard,
17+
types::AsNumber,
1818
types::{
1919
AsSequence, Comparable, Constructor, Hashable, Initializer, IterNext, IterNextIterable,
2020
Iterable, PyComparisonOp, Unconstructible, Unhashable,
@@ -23,6 +23,7 @@ use crate::{
2323
vm::VirtualMachine,
2424
AsObject, Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, TryFromObject,
2525
};
26+
use once_cell::sync::Lazy;
2627
use std::{fmt, ops::Deref};
2728

2829
pub type SetContentType = dictdatatype::Dict<()>;
@@ -488,7 +489,15 @@ fn reduce_set(
488489
}
489490

490491
#[pyclass(
491-
with(Constructor, Initializer, AsSequence, Hashable, Comparable, Iterable),
492+
with(
493+
Constructor,
494+
Initializer,
495+
AsSequence,
496+
Hashable,
497+
Comparable,
498+
Iterable,
499+
AsNumber
500+
),
492501
flags(BASETYPE)
493502
)]
494503
impl PySet {
@@ -792,6 +801,67 @@ impl Iterable for PySet {
792801
}
793802
}
794803

804+
impl AsNumber for PySet {
805+
fn as_number() -> &'static PyNumberMethods {
806+
static AS_NUMBER: Lazy<PyNumberMethods> = Lazy::new(|| PyNumberMethods {
807+
subtract: atomic_func!(|number, other, vm| {
808+
PySet::number_downcast(number)
809+
.sub(other.to_owned(), vm)
810+
.to_pyresult(vm)
811+
}),
812+
and: atomic_func!(|number, other, vm| {
813+
PySet::number_downcast(number)
814+
.and(other.to_owned(), vm)
815+
.to_pyresult(vm)
816+
}),
817+
xor: atomic_func!(|number, other, vm| {
818+
PySet::number_downcast(number)
819+
.xor(other.to_owned(), vm)
820+
.to_pyresult(vm)
821+
}),
822+
or: atomic_func!(|number, other, vm| {
823+
PySet::number_downcast(number)
824+
.or(other.to_owned(), vm)
825+
.to_pyresult(vm)
826+
}),
827+
inplace_subtract: atomic_func!(|number, other, vm| {
828+
PySet::isub(
829+
PySet::number_downcast(number).to_owned(),
830+
AnySet::try_from_object(vm, other.to_owned())?,
831+
vm,
832+
)
833+
.to_pyresult(vm)
834+
}),
835+
inplace_and: atomic_func!(|number, other, vm| {
836+
PySet::iand(
837+
PySet::number_downcast(number).to_owned(),
838+
AnySet::try_from_object(vm, other.to_owned())?,
839+
vm,
840+
)
841+
.to_pyresult(vm)
842+
}),
843+
inplace_xor: atomic_func!(|number, other, vm| {
844+
PySet::ixor(
845+
PySet::number_downcast(number).to_owned(),
846+
AnySet::try_from_object(vm, other.to_owned())?,
847+
vm,
848+
)
849+
.to_pyresult(vm)
850+
}),
851+
inplace_or: atomic_func!(|number, other, vm| {
852+
PySet::ior(
853+
PySet::number_downcast(number).to_owned(),
854+
AnySet::try_from_object(vm, other.to_owned())?,
855+
vm,
856+
)
857+
.to_pyresult(vm)
858+
}),
859+
..PyNumberMethods::NOT_IMPLEMENTED
860+
});
861+
&AS_NUMBER
862+
}
863+
}
864+
795865
impl Constructor for PyFrozenSet {
796866
type Args = OptionalArg<PyObjectRef>;
797867

@@ -822,7 +892,7 @@ impl Constructor for PyFrozenSet {
822892

823893
#[pyclass(
824894
flags(BASETYPE),
825-
with(Constructor, AsSequence, Hashable, Comparable, Iterable)
895+
with(Constructor, AsSequence, Hashable, Comparable, Iterable, AsNumber)
826896
)]
827897
impl PyFrozenSet {
828898
#[pymethod(magic)]
@@ -1019,6 +1089,35 @@ impl Iterable for PyFrozenSet {
10191089
}
10201090
}
10211091

1092+
impl AsNumber for PyFrozenSet {
1093+
fn as_number() -> &'static PyNumberMethods {
1094+
static AS_NUMBER: Lazy<PyNumberMethods> = Lazy::new(|| PyNumberMethods {
1095+
subtract: atomic_func!(|number, other, vm| {
1096+
PyFrozenSet::number_downcast(number)
1097+
.sub(other.to_owned(), vm)
1098+
.to_pyresult(vm)
1099+
}),
1100+
and: atomic_func!(|number, other, vm| {
1101+
PyFrozenSet::number_downcast(number)
1102+
.and(other.to_owned(), vm)
1103+
.to_pyresult(vm)
1104+
}),
1105+
xor: atomic_func!(|number, other, vm| {
1106+
PyFrozenSet::number_downcast(number)
1107+
.xor(other.to_owned(), vm)
1108+
.to_pyresult(vm)
1109+
}),
1110+
or: atomic_func!(|number, other, vm| {
1111+
PyFrozenSet::number_downcast(number)
1112+
.or(other.to_owned(), vm)
1113+
.to_pyresult(vm)
1114+
}),
1115+
..PyNumberMethods::NOT_IMPLEMENTED
1116+
});
1117+
&AS_NUMBER
1118+
}
1119+
}
1120+
10221121
struct AnySet {
10231122
object: PyObjectRef,
10241123
}

‎vm/src/builtins/type.rs

+22-7
Original file line numberDiff line numberDiff line change
@@ -2,12 +2,8 @@ use super::{
22
mappingproxy::PyMappingProxy, object, union_, PyClassMethod, PyDictRef, PyList, PyStaticMethod,
33
PyStr, PyStrInterned, PyStrRef, PyTuple, PyTupleRef, PyWeak,
44
};
5-
use crate::common::{
6-
ascii,
7-
borrow::BorrowedValue,
8-
lock::{PyRwLock, PyRwLockReadGuard},
9-
};
105
use crate::{
6+
atomic_func,
117
builtins::{
128
descriptor::{
139
DescrObject, MemberDef, MemberDescrObject, MemberGetter, MemberKind, MemberSetter,
@@ -17,15 +13,22 @@ use crate::{
1713
PyBaseExceptionRef,
1814
},
1915
class::{PyClassImpl, StaticType},
20-
convert::ToPyObject,
16+
common::{
17+
ascii,
18+
borrow::BorrowedValue,
19+
lock::{PyRwLock, PyRwLockReadGuard},
20+
},
21+
convert::{ToPyObject, ToPyResult},
2122
function::{FuncArgs, KwArgs, OptionalArg, PySetterValue},
2223
identifier,
2324
protocol::{PyIterReturn, PyMappingMethods, PyNumberMethods, PySequenceMethods},
25+
types::AsNumber,
2426
types::{Callable, GetAttr, PyTypeFlags, PyTypeSlots, SetAttr},
2527
AsObject, Context, Py, PyObjectRef, PyPayload, PyRef, PyResult, TryFromObject, VirtualMachine,
2628
};
2729
use indexmap::{map::Entry, IndexMap};
2830
use itertools::Itertools;
31+
use once_cell::sync::Lazy;
2932
use std::{borrow::Borrow, collections::HashSet, fmt, ops::Deref, pin::Pin, ptr::NonNull};
3033

3134
#[pyclass(module = false, name = "type")]
@@ -371,7 +374,7 @@ impl Py<PyType> {
371374
}
372375
}
373376

374-
#[pyclass(with(GetAttr, SetAttr, Callable), flags(BASETYPE))]
377+
#[pyclass(with(GetAttr, SetAttr, Callable, AsNumber), flags(BASETYPE))]
375378
impl PyType {
376379
// bound method for every type
377380
pub(crate) fn __new__(zelf: PyRef<Self>, args: FuncArgs, vm: &VirtualMachine) -> PyResult {
@@ -1058,6 +1061,18 @@ impl Callable for PyType {
10581061
}
10591062
}
10601063

1064+
impl AsNumber for PyType {
1065+
fn as_number() -> &'static PyNumberMethods {
1066+
static AS_NUMBER: Lazy<PyNumberMethods> = Lazy::new(|| PyNumberMethods {
1067+
or: atomic_func!(|num, other, vm| {
1068+
or_(num.obj.to_owned(), other.to_owned(), vm).to_pyresult(vm)
1069+
}),
1070+
..PyNumberMethods::NOT_IMPLEMENTED
1071+
});
1072+
&AS_NUMBER
1073+
}
1074+
}
1075+
10611076
fn find_base_dict_descr(cls: &Py<PyType>, vm: &VirtualMachine) -> Option<PyObjectRef> {
10621077
cls.iter_base_chain().skip(1).find_map(|cls| {
10631078
// TODO: should actually be some translation of:

0 commit comments

Comments
 (0)