Skip to content

Commit 7310fd4

Browse files
dpgeorgedpgeorge
committed
py: Consolidate min/max functions into one, and add key= argument.
Addresses issue adafruit#811.
1 parent 1d8a064 commit 7310fd4

2 files changed

Lines changed: 37 additions & 40 deletions

File tree

‎py/builtin.c‎

Lines changed: 27 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -284,63 +284,50 @@ STATIC mp_obj_t mp_builtin_iter(mp_obj_t o_in) {
284284

285285
MP_DEFINE_CONST_FUN_OBJ_1(mp_builtin_iter_obj, mp_builtin_iter);
286286

287-
STATIC mp_obj_t mp_builtin_max(uint n_args, const mp_obj_t *args) {
287+
STATIC mp_obj_t mp_builtin_min_max(uint n_args, const mp_obj_t *args, mp_map_t *kwargs, int op) {
288+
mp_map_elem_t *key_elem = mp_map_lookup(kwargs, MP_OBJ_NEW_QSTR(MP_QSTR_key), MP_MAP_LOOKUP);
289+
mp_obj_t key_fn = key_elem == NULL ? MP_OBJ_NULL : key_elem->value;
288290
if (n_args == 1) {
289291
// given an iterable
290292
mp_obj_t iterable = mp_getiter(args[0]);
291-
mp_obj_t max_obj = NULL;
293+
mp_obj_t best_key = MP_OBJ_NULL;
294+
mp_obj_t best_obj = MP_OBJ_NULL;
292295
mp_obj_t item;
293296
while ((item = mp_iternext(iterable)) != MP_OBJ_STOP_ITERATION) {
294-
if (max_obj == NULL || (mp_binary_op(MP_BINARY_OP_LESS, max_obj, item) == mp_const_true)) {
295-
max_obj = item;
297+
mp_obj_t key = key_fn == MP_OBJ_NULL ? item : mp_call_function_1(key_fn, item);
298+
if (best_obj == MP_OBJ_NULL || (mp_binary_op(op, key, best_key) == mp_const_true)) {
299+
best_key = key;
300+
best_obj = item;
296301
}
297302
}
298-
if (max_obj == NULL) {
299-
nlr_raise(mp_obj_new_exception_msg(&mp_type_ValueError, "max() arg is an empty sequence"));
303+
if (best_obj == MP_OBJ_NULL) {
304+
nlr_raise(mp_obj_new_exception_msg(&mp_type_ValueError, "arg is an empty sequence"));
300305
}
301-
return max_obj;
306+
return best_obj;
302307
} else {
303308
// given many args
304-
mp_obj_t max_obj = args[0];
305-
for (int i = 1; i < n_args; i++) {
306-
if (mp_binary_op(MP_BINARY_OP_LESS, max_obj, args[i]) == mp_const_true) {
307-
max_obj = args[i];
309+
mp_obj_t best_key = MP_OBJ_NULL;
310+
mp_obj_t best_obj = MP_OBJ_NULL;
311+
for (mp_uint_t i = 0; i < n_args; i++) {
312+
mp_obj_t key = key_fn == MP_OBJ_NULL ? args[i] : mp_call_function_1(key_fn, args[i]);
313+
if (best_obj == MP_OBJ_NULL || (mp_binary_op(op, key, best_key) == mp_const_true)) {
314+
best_key = key;
315+
best_obj = args[i];
308316
}
309317
}
310-
return max_obj;
318+
return best_obj;
311319
}
312320
}
313321

314-
MP_DEFINE_CONST_FUN_OBJ_VAR(mp_builtin_max_obj, 1, mp_builtin_max);
315-
316-
STATIC mp_obj_t mp_builtin_min(uint n_args, const mp_obj_t *args) {
317-
if (n_args == 1) {
318-
// given an iterable
319-
mp_obj_t iterable = mp_getiter(args[0]);
320-
mp_obj_t min_obj = NULL;
321-
mp_obj_t item;
322-
while ((item = mp_iternext(iterable)) != MP_OBJ_STOP_ITERATION) {
323-
if (min_obj == NULL || (mp_binary_op(MP_BINARY_OP_LESS, item, min_obj) == mp_const_true)) {
324-
min_obj = item;
325-
}
326-
}
327-
if (min_obj == NULL) {
328-
nlr_raise(mp_obj_new_exception_msg(&mp_type_ValueError, "min() arg is an empty sequence"));
329-
}
330-
return min_obj;
331-
} else {
332-
// given many args
333-
mp_obj_t min_obj = args[0];
334-
for (int i = 1; i < n_args; i++) {
335-
if (mp_binary_op(MP_BINARY_OP_LESS, args[i], min_obj) == mp_const_true) {
336-
min_obj = args[i];
337-
}
338-
}
339-
return min_obj;
340-
}
322+
STATIC mp_obj_t mp_builtin_max(uint n_args, const mp_obj_t *args, mp_map_t *kwargs) {
323+
return mp_builtin_min_max(n_args, args, kwargs, MP_BINARY_OP_MORE);
341324
}
325+
MP_DEFINE_CONST_FUN_OBJ_KW(mp_builtin_max_obj, 1, mp_builtin_max);
342326

343-
MP_DEFINE_CONST_FUN_OBJ_VAR(mp_builtin_min_obj, 1, mp_builtin_min);
327+
STATIC mp_obj_t mp_builtin_min(uint n_args, const mp_obj_t *args, mp_map_t *kwargs) {
328+
return mp_builtin_min_max(n_args, args, kwargs, MP_BINARY_OP_LESS);
329+
}
330+
MP_DEFINE_CONST_FUN_OBJ_KW(mp_builtin_min_obj, 1, mp_builtin_min);
344331

345332
STATIC mp_obj_t mp_builtin_next(mp_obj_t o) {
346333
mp_obj_t ret = mp_iternext_allow_raise(o);

‎tests/basics/builtin_minmax.py‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,3 +13,13 @@
1313
print(min([1,2,4,0,-1,2]))
1414
print(max([1,2,4,0,-1,2]))
1515

16+
# test with key function
17+
lst = [2, 1, 3, 4]
18+
print(min(lst, key=lambda x:x))
19+
print(min(lst, key=lambda x:-x))
20+
print(min(1, 2, 3, 4, key=lambda x:-x))
21+
print(min(4, 3, 2, 1, key=lambda x:-x))
22+
print(max(lst, key=lambda x:x))
23+
print(max(lst, key=lambda x:-x))
24+
print(max(1, 2, 3, 4, key=lambda x:-x))
25+
print(max(4, 3, 2, 1, key=lambda x:-x))

0 commit comments

Comments
 (0)