Skip to content

Commit 1d6fc94

Browse files
dpgeorgedpgeorge
committed
Implement framework for class-defined built-in operators.
Now working for class-defined methods: __getitem__, __setitem__, __add__, __sub__. Easy to add others.
1 parent 0c4e909 commit 1d6fc94

3 files changed

Lines changed: 100 additions & 0 deletions

File tree

‎py/objtype.c‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
#include "mpqstr.h"
1010
#include "obj.h"
1111
#include "map.h"
12+
#include "runtime0.h"
1213
#include "runtime.h"
1314

1415
/******************************************************************************/
@@ -105,6 +106,64 @@ static mp_obj_t class_make_new(mp_obj_t self_in, uint n_args, uint n_kw, const m
105106
return o;
106107
}
107108

109+
// TODO somehow replace const char * with a qstr
110+
static const char *binary_op_method_name[] = {
111+
[RT_BINARY_OP_SUBSCR] = "__getitem__",
112+
/*
113+
RT_BINARY_OP_OR,
114+
RT_BINARY_OP_XOR,
115+
RT_BINARY_OP_AND,
116+
RT_BINARY_OP_LSHIFT,
117+
RT_BINARY_OP_RSHIFT,
118+
*/
119+
[RT_BINARY_OP_ADD] = "__add__",
120+
[RT_BINARY_OP_SUBTRACT] = "__sub__",
121+
/*
122+
RT_BINARY_OP_MULTIPLY,
123+
RT_BINARY_OP_FLOOR_DIVIDE,
124+
RT_BINARY_OP_TRUE_DIVIDE,
125+
RT_BINARY_OP_MODULO,
126+
RT_BINARY_OP_POWER,
127+
RT_BINARY_OP_INPLACE_OR,
128+
RT_BINARY_OP_INPLACE_XOR,
129+
RT_BINARY_OP_INPLACE_AND,
130+
RT_BINARY_OP_INPLACE_LSHIFT,
131+
RT_BINARY_OP_INPLACE_RSHIFT,
132+
RT_BINARY_OP_INPLACE_ADD,
133+
RT_BINARY_OP_INPLACE_SUBTRACT,
134+
RT_BINARY_OP_INPLACE_MULTIPLY,
135+
RT_BINARY_OP_INPLACE_FLOOR_DIVIDE,
136+
RT_BINARY_OP_INPLACE_TRUE_DIVIDE,
137+
RT_BINARY_OP_INPLACE_MODULO,
138+
RT_BINARY_OP_INPLACE_POWER,
139+
RT_COMPARE_OP_LESS,
140+
RT_COMPARE_OP_MORE,
141+
RT_COMPARE_OP_EQUAL,
142+
RT_COMPARE_OP_LESS_EQUAL,
143+
RT_COMPARE_OP_MORE_EQUAL,
144+
RT_COMPARE_OP_NOT_EQUAL,
145+
RT_COMPARE_OP_IN,
146+
RT_COMPARE_OP_NOT_IN,
147+
RT_COMPARE_OP_IS,
148+
RT_COMPARE_OP_IS_NOT,
149+
*/
150+
[RT_COMPARE_OP_EXCEPTION_MATCH] = "__not_implemented__",
151+
};
152+
153+
static mp_obj_t class_binary_op(int op, mp_obj_t lhs_in, mp_obj_t rhs_in) {
154+
mp_obj_class_t *lhs = lhs_in;
155+
const char *op_name = binary_op_method_name[op];
156+
if (op_name == NULL) {
157+
return MP_OBJ_NULL;
158+
}
159+
mp_map_elem_t *elem = mp_obj_class_lookup(lhs->base.type, qstr_from_str_static(op_name), MP_MAP_LOOKUP);
160+
if (elem != NULL) {
161+
return rt_call_function_2(elem->value, lhs_in, rhs_in);
162+
} else {
163+
return MP_OBJ_NULL;
164+
}
165+
}
166+
108167
static void class_load_attr(mp_obj_t self_in, qstr attr, mp_obj_t *dest) {
109168
// logic: look in obj members then class locals (TODO check this against CPython)
110169
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) {
141200
return true;
142201
}
143202

203+
bool class_store_item(mp_obj_t self_in, mp_obj_t index, mp_obj_t value) {
204+
mp_obj_class_t *self = self_in;
205+
mp_map_elem_t *elem = mp_obj_class_lookup(self->base.type, qstr_from_str_static("__setitem__"), MP_MAP_LOOKUP);
206+
if (elem != NULL) {
207+
mp_obj_t args[3] = {self_in, index, value};
208+
return rt_call_function_n_kw(elem->value, 3, 0, args);
209+
} else {
210+
return MP_OBJ_NULL;
211+
}
212+
}
213+
144214
/******************************************************************************/
145215
// type object
146216
// - 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)
255325
o->name = qstr_str(name);
256326
o->print = class_print;
257327
o->make_new = class_make_new;
328+
o->binary_op = class_binary_op;
258329
o->load_attr = class_load_attr;
259330
o->store_attr = class_store_attr;
331+
o->store_item = class_store_item;
260332
o->bases_tuple = bases_tuple;
261333
o->locals_dict = locals_dict;
262334
return o;

‎tests/basics/tests/class_item.py‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
# test class with __getitem__ and __setitem__ methods
2+
3+
class C:
4+
def __getitem__(self, item):
5+
print('get', item)
6+
return 'item'
7+
8+
def __setitem__(self, item, value):
9+
print('set', item, value)
10+
11+
c = C()
12+
print(c[1])
13+
c[1] = 2

‎tests/basics/tests/class_number.py‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
# test class with __add__ and __sub__ methods
2+
3+
class C:
4+
def __init__(self, value):
5+
self.value = value
6+
7+
def __add__(self, rhs):
8+
print(self.value, '+', rhs)
9+
10+
def __sub__(self, rhs):
11+
print(self.value, '-', rhs)
12+
13+
c = C(0)
14+
c + 1
15+
c - 2

0 commit comments

Comments
 (0)