Skip to content

Commit 032129f

Browse files
committed
Implemented set.difference and set.difference_update
1 parent 2a24172 commit 032129f

2 files changed

Lines changed: 62 additions & 0 deletions

File tree

py/objset.c

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -137,6 +137,45 @@ static mp_obj_t set_discard(mp_obj_t self_in, mp_obj_t item) {
137137
}
138138
static MP_DEFINE_CONST_FUN_OBJ_2(set_discard_obj, set_discard);
139139

140+
static mp_obj_t set_diff_int(int n_args, const mp_obj_t *args, bool update) {
141+
assert(n_args > 0);
142+
assert(MP_OBJ_IS_TYPE(args[0], &set_type));
143+
mp_obj_set_t *self;
144+
if (update) {
145+
self = args[0];
146+
} else {
147+
self = set_copy(args[0]);
148+
}
149+
150+
151+
for (int i = 1; i < n_args; i++) {
152+
mp_obj_t other = args[i];
153+
if (self == other) {
154+
set_clear(self);
155+
} else {
156+
mp_obj_t iter = rt_getiter(other);
157+
mp_obj_t next;
158+
while ((next = rt_iternext(iter)) != mp_const_stop_iteration) {
159+
set_discard(self, next);
160+
}
161+
}
162+
}
163+
164+
return self;
165+
}
166+
167+
static mp_obj_t set_diff(int n_args, const mp_obj_t *args) {
168+
return set_diff_int(n_args, args, false);
169+
}
170+
static MP_DEFINE_CONST_FUN_OBJ_VAR(set_diff_obj, 1, set_diff);
171+
172+
static mp_obj_t set_diff_update(int n_args, const mp_obj_t *args) {
173+
set_diff_int(n_args, args, true);
174+
return mp_const_none;
175+
}
176+
static MP_DEFINE_CONST_FUN_OBJ_VAR(set_diff_update_obj, 1, set_diff_update);
177+
178+
140179
/******************************************************************************/
141180
/* set constructors & public C API */
142181

@@ -146,6 +185,8 @@ static const mp_method_t set_type_methods[] = {
146185
{ "clear", &set_clear_obj },
147186
{ "copy", &set_copy_obj },
148187
{ "discard", &set_discard_obj },
188+
{ "difference", &set_diff_obj },
189+
{ "difference_update", &set_diff_update_obj },
149190
{ NULL, NULL }, // end-of-list sentinel
150191
};
151192

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
def report(s):
2+
l = list(s)
3+
l.sort()
4+
print(l)
5+
6+
l = [1, 2, 3, 4]
7+
s = set(l)
8+
outs = [s.difference(),
9+
s.difference({1}),
10+
s.difference({1}, [1, 2]),
11+
s.difference({1}, {1, 2}, {2, 3})]
12+
for out in outs:
13+
report(out)
14+
15+
s = set(l)
16+
print(s.difference_update())
17+
report(s)
18+
print(s.difference_update({1}))
19+
report(s)
20+
print(s.difference_update({1}, [2]))
21+
report(s)

0 commit comments

Comments
 (0)