Skip to content

Commit 66a5bf6

Browse files
dpgeorgedpgeorge
committed
Merge pull request adafruit#142 from chipaca/containment
Implemented support for `in` and `not in` operators.
2 parents 0f59203 + f5a0a7d commit 66a5bf6

6 files changed

Lines changed: 100 additions & 6 deletions

File tree

‎py/objdict.c‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,12 @@ static mp_obj_t dict_binary_op(int op, mp_obj_t lhs_in, mp_obj_t rhs_in) {
5757
return elem->value;
5858
}
5959
}
60+
case RT_COMPARE_OP_IN:
61+
case RT_COMPARE_OP_NOT_IN:
62+
{
63+
mp_map_elem_t *elem = mp_map_lookup(&o->map, rhs_in, MP_MAP_LOOKUP);
64+
return MP_BOOL((op == RT_COMPARE_OP_IN) ^ (elem == NULL));
65+
}
6066
default:
6167
// op not supported
6268
return NULL;
@@ -362,10 +368,20 @@ static void dict_view_print(void (*print)(void *env, const char *fmt, ...), void
362368
print(env, "])");
363369
}
364370

371+
static mp_obj_t dict_view_binary_op(int op, mp_obj_t lhs_in, mp_obj_t rhs_in) {
372+
/* only supported for the 'keys' kind until sets and dicts are refactored */
373+
mp_obj_dict_view_t *o = lhs_in;
374+
if (o->kind != MP_DICT_VIEW_KEYS) return NULL;
375+
if (op != RT_COMPARE_OP_IN && op != RT_COMPARE_OP_NOT_IN) return NULL;
376+
return dict_binary_op(op, o->dict, rhs_in);
377+
}
378+
379+
365380
static const mp_obj_type_t dict_view_type = {
366381
{ &mp_const_type },
367382
"dict_view",
368383
.print = dict_view_print,
384+
.binary_op = dict_view_binary_op,
369385
.getiter = dict_view_getiter,
370386
};
371387

‎py/objset.c‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@ void set_print(void (*print)(void *env, const char *fmt, ...), void *env, mp_obj
4545
print(env, "}");
4646
}
4747

48+
4849
static mp_obj_t set_make_new(mp_obj_t type_in, int n_args, const mp_obj_t *args) {
4950
switch (n_args) {
5051
case 0:
@@ -405,6 +406,13 @@ static mp_obj_t set_binary_op(int op, mp_obj_t lhs, mp_obj_t rhs) {
405406
return set_issuperset(lhs, rhs);
406407
case RT_COMPARE_OP_NOT_EQUAL:
407408
return MP_BOOL(set_equal(lhs, rhs) == mp_const_false);
409+
case RT_COMPARE_OP_IN:
410+
case RT_COMPARE_OP_NOT_IN:
411+
{
412+
mp_obj_set_t *o = lhs;
413+
mp_obj_t elem = mp_set_lookup(&o->set, rhs, MP_MAP_LOOKUP);
414+
return MP_BOOL((op == RT_COMPARE_OP_IN) ^ (elem == NULL));
415+
}
408416
default:
409417
// op not supported
410418
return NULL;

‎py/objstr.c‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,15 @@ mp_obj_t str_binary_op(int op, mp_obj_t lhs_in, mp_obj_t rhs_in) {
8585
return mp_obj_new_str(qstr_from_str_take(val, alloc_len));
8686
}
8787
break;
88+
case RT_COMPARE_OP_IN:
89+
case RT_COMPARE_OP_NOT_IN:
90+
/* NOTE `a in b` is `b.__contains__(a)` */
91+
if (MP_OBJ_IS_TYPE(rhs_in, &str_type)) {
92+
const char *rhs_str = qstr_str(((mp_obj_str_t*)rhs_in)->qstr);
93+
/* FIXME \0 in strs */
94+
return MP_BOOL((op == RT_COMPARE_OP_IN) ^ (strstr(lhs_str, rhs_str) == NULL));
95+
}
96+
break;
8897
}
8998

9099
return MP_OBJ_NULL; // op not supported

‎py/runtime.c‎

Lines changed: 41 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -568,22 +568,57 @@ mp_obj_t rt_binary_op(int op, mp_obj_t lhs, mp_obj_t rhs) {
568568
} else if (MP_OBJ_IS_TYPE(rhs, &complex_type)) {
569569
return mp_obj_complex_binary_op(op, lhs_val, 0, rhs);
570570
}
571-
} else {
572-
if (MP_OBJ_IS_OBJ(lhs)) {
573-
mp_obj_base_t *o = lhs;
571+
}
572+
573+
/* deal with `in` and `not in`
574+
*
575+
* NOTE `a in b` is `b.__contains__(a)`, hence why the generic dispatch
576+
* needs to go below
577+
*/
578+
if (op == RT_COMPARE_OP_IN || op == RT_COMPARE_OP_NOT_IN) {
579+
if (!MP_OBJ_IS_SMALL_INT(rhs)) {
580+
mp_obj_base_t *o = rhs;
574581
if (o->type->binary_op != NULL) {
575-
mp_obj_t result = o->type->binary_op(op, lhs, rhs);
576-
if (result != NULL) {
577-
return result;
582+
mp_obj_t res = o->type->binary_op(op, rhs, lhs);
583+
if (res != NULL) {
584+
return res;
578585
}
579586
}
587+
if (o->type->getiter != NULL) {
588+
/* second attempt, walk the iterator */
589+
mp_obj_t next = NULL;
590+
mp_obj_t iter = rt_getiter(rhs);
591+
while ((next = rt_iternext(iter)) != mp_const_stop_iteration) {
592+
if (mp_obj_equal(next, lhs)) {
593+
return MP_BOOL(op == RT_COMPARE_OP_IN);
594+
}
595+
}
596+
return MP_BOOL(op != RT_COMPARE_OP_IN);
597+
}
598+
}
599+
600+
nlr_jump(mp_obj_new_exception_msg_varg(
601+
MP_QSTR_TypeError, "'%s' object is not iterable",
602+
mp_obj_get_type_str(rhs)));
603+
return mp_const_none;
604+
}
605+
606+
if (MP_OBJ_IS_OBJ(lhs)) {
607+
mp_obj_base_t *o = lhs;
608+
if (o->type->binary_op != NULL) {
609+
mp_obj_t result = o->type->binary_op(op, lhs, rhs);
610+
if (result != NULL) {
611+
return result;
612+
}
580613
}
614+
// TODO implement dispatch for reverse binary ops
581615
}
582616

583617
// TODO specify in error message what the operator is
584618
nlr_jump(mp_obj_new_exception_msg_varg(MP_QSTR_TypeError,
585619
"unsupported operand types for binary operator: '%s', '%s'",
586620
mp_obj_get_type_str(lhs), mp_obj_get_type_str(rhs)));
621+
return mp_const_none;
587622
}
588623

589624
mp_obj_t rt_make_function_from_id(int unique_code_id) {

‎tests/basics/run-tests‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,4 +42,7 @@ echo "$numpassed tests passed"
4242
if [[ $numfailed != 0 ]]
4343
then
4444
echo "$numfailed tests failed -$namefailed"
45+
exit 1
46+
else
47+
exit 0
4548
fi

‎tests/basics/tests/containment.py‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
1+
for i in 1, 2:
2+
for o in {1:2}, {1}, {1:2}.keys():
3+
print("{} in {}: {}".format(i, o, i in o))
4+
print("{} not in {}: {}".format(i, o, i not in o))
5+
6+
haystack = "supercalifragilistc"
7+
for needle in (haystack[i:] for i in range(len(haystack))):
8+
print(needle, "in", haystack, "::", needle in haystack)
9+
print(needle, "not in", haystack, "::", needle not in haystack)
10+
print(haystack, "in", needle, "::", haystack in needle)
11+
print(haystack, "not in", needle, "::", haystack not in needle)
12+
for needle in (haystack[:i+1] for i in range(len(haystack))):
13+
print(needle, "in", haystack, "::", needle in haystack)
14+
print(needle, "not in", haystack, "::", needle not in haystack)
15+
print(haystack, "in", needle, "::", haystack in needle)
16+
print(haystack, "not in", needle, "::", haystack not in needle)
17+
18+
# until here, the tests would work without the 'second attempt' iteration thing.
19+
20+
for i in 1, 2:
21+
for o in [], [1], [1, 2]:
22+
print("{} in {}: {}".format(i, o, i in o))
23+
print("{} not in {}: {}".format(i, o, i not in o))

0 commit comments

Comments
 (0)