Skip to content

Commit

Permalink
Implemented __and__, __or__ and __sub__
Browse files Browse the repository at this point in the history
  • Loading branch information
@azmeuk
azmeuk committed Jun 2, 2020
1 parent 078ba60 commit 5f15484
Show file tree
Hide file tree
Showing 9 changed files with 128 additions and 5 deletions.
13 changes: 11 additions & 2 deletions BTrees/BTreeTemplate.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2425,7 +2425,7 @@ BTree_nonzero(BTree *self)

static PyNumberMethods BTree_as_number_for_nonzero = {
0, /* nb_add */
0, /* nb_subtract */
bucket_sub, /* nb_subtract */
0, /* nb_multiply */
#ifndef PY3K
0, /* nb_divide */
Expand All @@ -2436,7 +2436,13 @@ static PyNumberMethods BTree_as_number_for_nonzero = {
0, /* nb_negative */
0, /* nb_positive */
0, /* nb_absolute */
(inquiry)BTree_nonzero /* nb_nonzero */
(inquiry)BTree_nonzero, /* nb_nonzero */
(unaryfunc)0, /* nb_invert */
(binaryfunc)0, /* nb_lshift */
(binaryfunc)0, /* nb_rshift */
bucket_and, /* nb_and */
(binaryfunc)0, /* nb_xor */
bucket_or, /* nb_or */
};

static PyTypeObject BTreeType = {
Expand All @@ -2459,6 +2465,9 @@ static PyTypeObject BTreeType = {
0, /* tp_getattro */
0, /* tp_setattro */
0, /* tp_as_buffer */
#ifndef PY3K
Py_TPFLAGS_CHECKTYPES |
#endif
Py_TPFLAGS_DEFAULT |
Py_TPFLAGS_HAVE_GC |
Py_TPFLAGS_BASETYPE, /* tp_flags */
Expand Down
49 changes: 48 additions & 1 deletion BTrees/BucketTemplate.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
****************************************************************************/

#include "SetOpTemplate.h"
#define BUCKETTEMPLATE_C "$Id$\n"

/* Use BUCKET_SEARCH to find the index at which a key belongs.
Expand Down Expand Up @@ -1367,6 +1368,26 @@ bucket_setstate(Bucket *self, PyObject *state)
return Py_None;
}

static PyObject *
bucket_sub(PyObject *self, PyObject *other)
{
PyObject *args = Py_BuildValue("OO", self, other);
return difference_m(NULL, args);
}

static PyObject *
bucket_or(PyObject *self, PyObject *other)
{
PyObject *args = Py_BuildValue("OO", self, other);
return union_m(NULL, args);
}

static PyObject *
bucket_and(PyObject *self, PyObject *other)
{
PyObject *args = Py_BuildValue("OO", self, other);
return intersection_m(NULL, args);
}

static PyObject *
bucket_setdefault(Bucket *self, PyObject *args)
Expand Down Expand Up @@ -1847,6 +1868,29 @@ static PySequenceMethods Bucket_as_sequence = {
0, /* sq_inplace_repeat */
};

static PyNumberMethods Bucket_as_number = {
(binaryfunc)0, /* nb_add */
bucket_sub, /* nb_subtract */
(binaryfunc)0, /* nb_multiply */
#ifndef PY3K
0, /* nb_divide */
#endif
(binaryfunc)0, /* nb_remainder */
(binaryfunc)0, /* nb_divmod */
(ternaryfunc)0, /* nb_power */
(unaryfunc)0, /* nb_negative */
(unaryfunc)0, /* nb_positive */
(unaryfunc)0, /* nb_absolute */
(inquiry)0, /* nb_bool */
(unaryfunc)0, /* nb_invert */
(binaryfunc)0, /* nb_lshift */
(binaryfunc)0, /* nb_rshift */
bucket_and, /* nb_and */
(binaryfunc)0, /* nb_xor */
bucket_or, /* nb_or */
};


static PyObject *
bucket_repr(Bucket *self)
{
Expand Down Expand Up @@ -1911,7 +1955,7 @@ static PyTypeObject BucketType = {
0, /* tp_setattr */
0, /* tp_compare */
(reprfunc)bucket_repr, /* tp_repr */
0, /* tp_as_number */
&Bucket_as_number, /* tp_as_number */
&Bucket_as_sequence, /* tp_as_sequence */
&Bucket_as_mapping, /* tp_as_mapping */
0, /* tp_hash */
Expand All @@ -1920,6 +1964,9 @@ static PyTypeObject BucketType = {
0, /* tp_getattro */
0, /* tp_setattro */
0, /* tp_as_buffer */
#ifndef PY3K
Py_TPFLAGS_CHECKTYPES |
#endif
Py_TPFLAGS_DEFAULT |
Py_TPFLAGS_HAVE_GC |
Py_TPFLAGS_BASETYPE, /* tp_flags */
Expand Down
1 change: 1 addition & 0 deletions BTrees/SetOpTemplate.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
Set operations
****************************************************************************/

#include "SetOpTemplate.h"
#define SETOPTEMPLATE_C "$Id$\n"

#ifdef KEY_CHECK
Expand Down
15 changes: 15 additions & 0 deletions BTrees/SetOpTemplate.h
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,15 @@
# ifndef SETOPTEMPLATE_H
# define SETOPTEMPLATE_H

#include "Python.h"

static PyObject *
union_m(PyObject *ignored, PyObject *args);

static PyObject *
intersection_m(PyObject *ignored, PyObject *args);

static PyObject *
difference_m(PyObject *ignored, PyObject *args);

# endif
27 changes: 26 additions & 1 deletion BTrees/SetTemplate.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -304,6 +304,28 @@ static PySequenceMethods set_as_sequence = {
0, /* sq_inplace_repeat */
};

static PyNumberMethods set_as_number = {
(binaryfunc)0, /* nb_add */
bucket_sub, /* nb_subtract */
(binaryfunc)0, /* nb_multiply */
#ifndef PY3K
0, /* nb_divide */
#endif
(binaryfunc)0, /* nb_remainder */
(binaryfunc)0, /* nb_divmod */
(ternaryfunc)0, /* nb_power */
(unaryfunc)0, /* nb_negative */
(unaryfunc)0, /* nb_positive */
(unaryfunc)0, /* nb_absolute */
(inquiry)0, /* nb_bool */
(unaryfunc)0, /* nb_invert */
(binaryfunc)0, /* nb_lshift */
(binaryfunc)0, /* nb_rshift */
bucket_and, /* nb_and */
(binaryfunc)0, /* nb_xor */
bucket_or, /* nb_or */
};

static PyTypeObject SetType = {
PyVarObject_HEAD_INIT(NULL, 0) /* PyPersist_Type */
MODULE_NAME MOD_NAME_PREFIX "Set", /* tp_name */
Expand All @@ -315,7 +337,7 @@ static PyTypeObject SetType = {
0, /* tp_setattr */
0, /* tp_compare */
(reprfunc)set_repr, /* tp_repr */
0, /* tp_as_number */
&set_as_number, /* tp_as_number */
&set_as_sequence, /* tp_as_sequence */
0, /* tp_as_mapping */
0, /* tp_hash */
Expand All @@ -324,6 +346,9 @@ static PyTypeObject SetType = {
0, /* tp_getattro */
0, /* tp_setattro */
0, /* tp_as_buffer */
#ifndef PY3K
Py_TPFLAGS_CHECKTYPES |
#endif
Py_TPFLAGS_DEFAULT |
Py_TPFLAGS_HAVE_GC |
Py_TPFLAGS_BASETYPE, /* tp_flags */
Expand Down
3 changes: 3 additions & 0 deletions BTrees/TreeSetTemplate.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -230,6 +230,9 @@ static PyTypeObject TreeSetType =
0, /* tp_getattro */
0, /* tp_setattro */
0, /* tp_as_buffer */
#ifndef PY3K
Py_TPFLAGS_CHECKTYPES |
#endif
Py_TPFLAGS_DEFAULT |
Py_TPFLAGS_HAVE_GC |
Py_TPFLAGS_BASETYPE, /* tp_flags */
Expand Down
20 changes: 20 additions & 0 deletions BTrees/_base.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -218,6 +218,16 @@ def _repr_helper(self, items):
name = name[:-2] if name.endswith("Py") else name
return "%s.%s(%r)" % (mod, name, items)

def __sub__(self, other):
return difference(self.__class__, self, other)

def __or__(self, other):
return union(self.__class__, self, other)

def __and__(self, other):
return intersection(self.__class__, self, other)


class _SetIteration(object):

__slots__ = ('to_iterate',
Expand Down Expand Up @@ -1143,6 +1153,16 @@ def __repr__(self):
r = r.replace('Py', '')
return r

def __sub__(self, other):
return difference(self.__class__, self, other)

def __or__(self, other):
return union(self.__class__, self, other)

def __and__(self, other):
return intersection(self.__class__, self, other)


def _get_simple_btree_bucket_state(state):
if state is None:
return state
Expand Down
3 changes: 3 additions & 0 deletions BTrees/tests/common.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2313,6 +2313,7 @@ def testUnion(self):
C = self.union(A, B)
self.assertTrue(not hasattr(C, "values"))
self.assertEqual(list(C), self._union(A, B))
self.assertEqual(set(A) | set(B), set(A | B))

def testIntersection(self):
inputs = self.As + self.Bs
Expand All @@ -2321,6 +2322,7 @@ def testIntersection(self):
C = self.intersection(A, B)
self.assertTrue(not hasattr(C, "values"))
self.assertEqual(list(C), self._intersection(A, B))
self.assertEqual(set(A) & set(B), set(A & B))

def testDifference(self):
inputs = self.As + self.Bs
Expand All @@ -2334,6 +2336,7 @@ def testDifference(self):
self.assertEqual(list(C.items()), want)
else:
self.assertEqual(list(C), want)
self.assertEqual(set(A) - set(B), set(A - B))

def testLargerInputs(self): # pylint:disable=too-many-locals
from BTrees.IIBTree import IISet # pylint:disable=no-name-in-module
Expand Down
2 changes: 1 addition & 1 deletion tox.ini
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@ extras =
commands =
# Temporary work around. Avoid zope.testrunner pending
# https://github.com/zopefoundation/zope.security/issues/71
python -m unittest discover -s BTrees -t .
python -m unittest discover -s BTrees -t . {posargs}
setenv =
PYTHONFAULTHANDLER=1
PYTHONDEVMODE=1
Expand Down

0 comments on commit 5f15484

Please sign in to comment.