Skip to content

Commit a41fe31

Browse files
chipacachipaca
committed
Added dict iterator.
1 parent 24507af commit a41fe31

2 files changed

Lines changed: 76 additions & 10 deletions

File tree

‎py/objdict.c‎

Lines changed: 73 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -17,20 +17,23 @@ typedef struct _mp_obj_dict_t {
1717
mp_map_t map;
1818
} mp_obj_dict_t;
1919

20+
static mp_obj_t mp_obj_new_dict_iterator(mp_obj_dict_t *dict, int cur);
21+
static mp_map_elem_t *dict_it_iternext_elem(mp_obj_t self_in);
22+
2023
static void dict_print(void (*print)(void *env, const char *fmt, ...), void *env, mp_obj_t self_in) {
2124
mp_obj_dict_t *self = self_in;
2225
bool first = true;
2326
print(env, "{");
24-
for (int i = 0; i < self->map.alloc; i++) {
25-
if (self->map.table[i].key != NULL) {
26-
if (!first) {
27-
print(env, ", ");
28-
}
29-
first = false;
30-
mp_obj_print_helper(print, env, self->map.table[i].key);
31-
print(env, ": ");
32-
mp_obj_print_helper(print, env, self->map.table[i].value);
27+
mp_obj_t *dict_iter = mp_obj_new_dict_iterator(self, 0);
28+
mp_map_elem_t *next = NULL;
29+
while ((next = dict_it_iternext_elem(dict_iter)) != NULL) {
30+
if (!first) {
31+
print(env, ", ");
3332
}
33+
first = false;
34+
mp_obj_print_helper(print, env, next->key);
35+
print(env, ": ");
36+
mp_obj_print_helper(print, env, next->value);
3437
}
3538
print(env, "}");
3639
}
@@ -60,13 +63,73 @@ static mp_obj_t dict_binary_op(int op, mp_obj_t lhs_in, mp_obj_t rhs_in) {
6063
}
6164
}
6265

66+
67+
/******************************************************************************/
68+
/* dict iterator */
69+
70+
typedef struct _mp_obj_dict_it_t {
71+
mp_obj_base_t base;
72+
mp_obj_dict_t *dict;
73+
machine_uint_t cur;
74+
} mp_obj_dict_it_t;
75+
76+
static mp_map_elem_t *dict_it_iternext_elem(mp_obj_t self_in) {
77+
mp_obj_dict_it_t *self = self_in;
78+
machine_uint_t max = self->dict->map.alloc;
79+
mp_map_elem_t *table = self->dict->map.table;
80+
81+
for (int i = self->cur; i < max; i++) {
82+
if (table[i].key != NULL) {
83+
self->cur = i + 1;
84+
return &(table[i]);
85+
}
86+
}
87+
88+
return NULL;
89+
}
90+
91+
mp_obj_t dict_it_iternext(mp_obj_t self_in) {
92+
mp_map_elem_t *next = dict_it_iternext_elem(self_in);
93+
94+
if (next != NULL) {
95+
return next->key;
96+
} else {
97+
return mp_const_stop_iteration;
98+
}
99+
}
100+
101+
static const mp_obj_type_t dict_it_type = {
102+
{ &mp_const_type },
103+
"dict_iterator",
104+
.iternext = dict_it_iternext,
105+
.methods = { { NULL, NULL }, },
106+
};
107+
108+
static mp_obj_t mp_obj_new_dict_iterator(mp_obj_dict_t *dict, int cur) {
109+
mp_obj_dict_it_t *o = m_new_obj(mp_obj_dict_it_t);
110+
o->base.type = &dict_it_type;
111+
o->dict = dict;
112+
o->cur = cur;
113+
return o;
114+
}
115+
116+
static mp_obj_t dict_getiter(mp_obj_t o_in) {
117+
return mp_obj_new_dict_iterator(o_in, 0);
118+
}
119+
120+
/******************************************************************************/
121+
/* dict methods */
122+
123+
/******************************************************************************/
124+
/* dict constructors & etc */
125+
63126
const mp_obj_type_t dict_type = {
64127
{ &mp_const_type },
65128
"dict",
66129
.print = dict_print,
67130
.make_new = dict_make_new,
68131
.binary_op = dict_binary_op,
69-
.getiter = NULL,
132+
.getiter = dict_getiter,
70133
.methods = {{NULL, NULL},},
71134
};
72135

‎tests/basics/tests/dict_iterator.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
d = {1: 2, 3: 4}
2+
for i in d:
3+
print(i, d[i])

0 commit comments

Comments
 (0)