diff options
| -rw-r--r-- | py/objtype.c | 72 | ||||
| -rw-r--r-- | tests/basics/tests/class_item.py | 13 | ||||
| -rw-r--r-- | tests/basics/tests/class_number.py | 15 |
3 files changed, 100 insertions, 0 deletions
diff --git a/py/objtype.c b/py/objtype.c index 2dcca6b9e1..3ae56eb708 100644 --- a/py/objtype.c +++ b/py/objtype.c @@ -9,6 +9,7 @@ #include "mpqstr.h" #include "obj.h" #include "map.h" +#include "runtime0.h" #include "runtime.h" /******************************************************************************/ @@ -105,6 +106,64 @@ static mp_obj_t class_make_new(mp_obj_t self_in, uint n_args, uint n_kw, const m return o; } +// TODO somehow replace const char * with a qstr +static const char *binary_op_method_name[] = { + [RT_BINARY_OP_SUBSCR] = "__getitem__", + /* + RT_BINARY_OP_OR, + RT_BINARY_OP_XOR, + RT_BINARY_OP_AND, + RT_BINARY_OP_LSHIFT, + RT_BINARY_OP_RSHIFT, + */ + [RT_BINARY_OP_ADD] = "__add__", + [RT_BINARY_OP_SUBTRACT] = "__sub__", + /* + RT_BINARY_OP_MULTIPLY, + RT_BINARY_OP_FLOOR_DIVIDE, + RT_BINARY_OP_TRUE_DIVIDE, + RT_BINARY_OP_MODULO, + RT_BINARY_OP_POWER, + RT_BINARY_OP_INPLACE_OR, + RT_BINARY_OP_INPLACE_XOR, + RT_BINARY_OP_INPLACE_AND, + RT_BINARY_OP_INPLACE_LSHIFT, + RT_BINARY_OP_INPLACE_RSHIFT, + RT_BINARY_OP_INPLACE_ADD, + RT_BINARY_OP_INPLACE_SUBTRACT, + RT_BINARY_OP_INPLACE_MULTIPLY, + RT_BINARY_OP_INPLACE_FLOOR_DIVIDE, + RT_BINARY_OP_INPLACE_TRUE_DIVIDE, + RT_BINARY_OP_INPLACE_MODULO, + RT_BINARY_OP_INPLACE_POWER, + RT_COMPARE_OP_LESS, + RT_COMPARE_OP_MORE, + RT_COMPARE_OP_EQUAL, + RT_COMPARE_OP_LESS_EQUAL, + RT_COMPARE_OP_MORE_EQUAL, + RT_COMPARE_OP_NOT_EQUAL, + RT_COMPARE_OP_IN, + RT_COMPARE_OP_NOT_IN, + RT_COMPARE_OP_IS, + RT_COMPARE_OP_IS_NOT, + */ + [RT_COMPARE_OP_EXCEPTION_MATCH] = "__not_implemented__", +}; + +static mp_obj_t class_binary_op(int op, mp_obj_t lhs_in, mp_obj_t rhs_in) { + mp_obj_class_t *lhs = lhs_in; + const char *op_name = binary_op_method_name[op]; + if (op_name == NULL) { + return MP_OBJ_NULL; + } + mp_map_elem_t *elem = mp_obj_class_lookup(lhs->base.type, qstr_from_str_static(op_name), MP_MAP_LOOKUP); + if (elem != NULL) { + return rt_call_function_2(elem->value, lhs_in, rhs_in); + } else { + return MP_OBJ_NULL; + } +} + static void class_load_attr(mp_obj_t self_in, qstr attr, mp_obj_t *dest) { // logic: look in obj members then class locals (TODO check this against CPython) mp_obj_class_t *self = self_in; @@ -141,6 +200,17 @@ static bool class_store_attr(mp_obj_t self_in, qstr attr, mp_obj_t value) { return true; } +bool class_store_item(mp_obj_t self_in, mp_obj_t index, mp_obj_t value) { + mp_obj_class_t *self = self_in; + mp_map_elem_t *elem = mp_obj_class_lookup(self->base.type, qstr_from_str_static("__setitem__"), MP_MAP_LOOKUP); + if (elem != NULL) { + mp_obj_t args[3] = {self_in, index, value}; + return rt_call_function_n_kw(elem->value, 3, 0, args); + } else { + return MP_OBJ_NULL; + } +} + /******************************************************************************/ // type object // - the struct is mp_obj_type_t and is defined in obj.h so const types can be made @@ -255,8 +325,10 @@ mp_obj_t mp_obj_new_type(qstr name, mp_obj_t bases_tuple, mp_obj_t locals_dict) o->name = qstr_str(name); o->print = class_print; o->make_new = class_make_new; + o->binary_op = class_binary_op; o->load_attr = class_load_attr; o->store_attr = class_store_attr; + o->store_item = class_store_item; o->bases_tuple = bases_tuple; o->locals_dict = locals_dict; return o; diff --git a/tests/basics/tests/class_item.py b/tests/basics/tests/class_item.py new file mode 100644 index 0000000000..6061f26075 --- /dev/null +++ b/tests/basics/tests/class_item.py @@ -0,0 +1,13 @@ +# test class with __getitem__ and __setitem__ methods + +class C: + def __getitem__(self, item): + print('get', item) + return 'item' + + def __setitem__(self, item, value): + print('set', item, value) + +c = C() +print(c[1]) +c[1] = 2 diff --git a/tests/basics/tests/class_number.py b/tests/basics/tests/class_number.py new file mode 100644 index 0000000000..e1dbf4a26c --- /dev/null +++ b/tests/basics/tests/class_number.py @@ -0,0 +1,15 @@ +# test class with __add__ and __sub__ methods + +class C: + def __init__(self, value): + self.value = value + + def __add__(self, rhs): + print(self.value, '+', rhs) + + def __sub__(self, rhs): + print(self.value, '-', rhs) + +c = C(0) +c + 1 +c - 2 |