diff --git a/include/typings/array2d.pyi b/include/typings/array2d.pyi index 3289d5a8..79debad9 100644 --- a/include/typings/array2d.pyi +++ b/include/typings/array2d.pyi @@ -35,6 +35,7 @@ class array2d_like[T]: def map[R](self, f: Callable[[T], R]) -> array2d[R]: ... def apply(self, f: Callable[[T], T]) -> None: ... + def zip_with[R, U](self, other: array2d_like[U], f: Callable[[T, U], R]) -> array2d[R]: ... def copy(self) -> 'array2d[T]': ... def tolist(self) -> list[list[T]]: ... @@ -50,6 +51,8 @@ class array2d_like[T]: @overload def __getitem__(self, index: tuple[slice, slice]) -> array2d_view[T]: ... @overload + def __getitem__(self, index: tuple[slice, int] | tuple[int, slice]) -> array2d_view[T]: ... + @overload def __getitem__(self, mask: array2d_like[bool]) -> list[T]: ... @overload def __setitem__(self, index: vec2i, value: T): ... @@ -58,6 +61,8 @@ class array2d_like[T]: @overload def __setitem__(self, index: tuple[slice, slice], value: T | 'array2d_like[T]'): ... @overload + def __setitem__(self, index: tuple[slice, int] | tuple[int, slice], value: T | 'array2d_like[T]'): ... + @overload def __setitem__(self, mask: array2d_like[bool], value: T): ... # algorithms diff --git a/src/interpreter/vm.c b/src/interpreter/vm.c index 2876d815..88fc7794 100644 --- a/src/interpreter/vm.c +++ b/src/interpreter/vm.c @@ -282,6 +282,18 @@ static void _clip_int(int* value, int min, int max) { } bool pk__parse_int_slice(py_Ref slice, int length, int* start, int* stop, int* step) { + if(py_isint(slice)) { + int index = py_toint(slice); + bool ok = pk__normalize_index(&index, length); + if(!ok) return false; + *start = index; + *stop = index + 1; + *step = 1; + return true; + } + + if(!py_istype(slice, tp_slice)) c11__abort("pk__parse_int_slice(): not a slice object"); + py_Ref s_start = py_getslot(slice, 0); py_Ref s_stop = py_getslot(slice, 1); py_Ref s_step = py_getslot(slice, 2); @@ -334,7 +346,7 @@ bool pk__parse_int_slice(py_Ref slice, int length, int* start, int* stop, int* s bool pk__normalize_index(int* index, int length) { if(*index < 0) *index += length; - if(*index < 0 || *index >= length) { return IndexError("%d not in [0, %d)", *index, length); } + if(*index < 0 || *index >= length) return IndexError("%d not in [0, %d)", *index, length); return true; } diff --git a/src/modules/array2d.c b/src/modules/array2d.c index b5b08850..4eb9247e 100644 --- a/src/modules/array2d.c +++ b/src/modules/array2d.c @@ -188,6 +188,41 @@ static bool array2d_like_apply(int argc, py_Ref argv) { return true; } +static bool _check_same_shape(int colA, int rowA, int colB, int rowB) { + if(colA != colB || rowA != rowB) { + const char* fmt = "expected the same shape: (%d, %d) != (%d, %d)"; + return ValueError(fmt, colA, rowA, colB, rowB); + } + return true; +} + +static bool _array2d_like_check_same_shape(c11_array2d_like* self, c11_array2d_like* other) { + return _check_same_shape(self->n_cols, self->n_rows, other->n_cols, other->n_rows); +} + +static bool array2d_like_zip_with(int argc, py_Ref argv) { + PY_CHECK_ARGC(3); + c11_array2d_like* self = py_touserdata(argv); + if(!py_checkinstance(py_arg(1), tp_array2d_like)) return false; + c11_array2d_like* other = py_touserdata(py_arg(1)); + py_Ref f = py_arg(2); + if(!_array2d_like_check_same_shape(self, other)) return false; + c11_array2d* res = py_newarray2d(py_pushtmp(), self->n_cols, self->n_rows); + for(int j = 0; j < self->n_rows; j++) { + for(int i = 0; i < self->n_cols; i++) { + py_push(f); + py_pushnil(); + py_push(self->f_get(self, i, j)); + py_push(other->f_get(other, i, j)); + if(!py_vectorcall(2, 0)) return false; + c11_array2d__set(res, i, j, py_retval()); + } + } + py_assign(py_retval(), py_peek(-1)); + py_pop(); + return true; +} + static bool array2d_like_copy(int argc, py_Ref argv) { // def copy(self) -> 'array2d': ... PY_CHECK_ARGC(1); @@ -217,18 +252,6 @@ static bool array2d_like_tolist(int argc, py_Ref argv) { return true; } -static bool _check_same_shape(int colA, int rowA, int colB, int rowB) { - if(colA != colB || rowA != rowB) { - const char* fmt = "expected the same shape: (%d, %d) != (%d, %d)"; - return ValueError(fmt, colA, rowA, colB, rowB); - } - return true; -} - -static bool _array2d_like_check_same_shape(c11_array2d_like* self, c11_array2d_like* other) { - return _check_same_shape(self->n_cols, self->n_rows, other->n_cols, other->n_rows); -} - static bool array2d_like__eq__(int argc, py_Ref argv) { PY_CHECK_ARGC(2); c11_array2d_like* self = py_touserdata(argv); @@ -418,7 +441,10 @@ static bool array2d_like__getitem__(int argc, py_Ref argv) { return _array2d_like_IndexError(self, col, row); } - if(py_istype(x, tp_slice) && py_istype(y, tp_slice)) { + bool _1 = py_istype(x, tp_slice) && py_istype(y, tp_slice); + bool _2 = py_istype(x, tp_int) && py_istype(y, tp_slice); + bool _3 = py_istype(x, tp_slice) && py_istype(y, tp_int); + if(_1 || _2 || _3) { HANDLE_SLICE(); return _array2d_view(py_retval(), argv, @@ -429,7 +455,7 @@ static bool array2d_like__getitem__(int argc, py_Ref argv) { slice_height); } - return TypeError("expected `tuple[int, int]` or `tuple[slice, slice]`"); + return TypeError("expected tuple[int, int] or tuple[slice, slice]"); } static bool array2d_like__setitem__(int argc, py_Ref argv) { @@ -480,7 +506,10 @@ static bool array2d_like__setitem__(int argc, py_Ref argv) { return _array2d_like_IndexError(self, col, row); } - if(py_istype(x, tp_slice) && py_istype(y, tp_slice)) { + bool _1 = py_istype(x, tp_slice) && py_istype(y, tp_slice); + bool _2 = py_istype(x, tp_int) && py_istype(y, tp_slice); + bool _3 = py_istype(x, tp_slice) && py_istype(y, tp_int); + if(_1 || _2 || _3) { HANDLE_SLICE(); if(py_isinstance(value, tp_array2d_like)) { c11_array2d_like* values = py_touserdata(value); @@ -505,7 +534,7 @@ static bool array2d_like__setitem__(int argc, py_Ref argv) { return true; } - return TypeError("expected `tuple[int, int]` or `tuple[slice, slice]"); + return TypeError("expected tuple[int, int] or tuple[slice, slice]"); } // count(self, value: T) -> int @@ -684,6 +713,7 @@ static void register_array2d_like(py_Ref mod) { py_bindmethod(type, "map", array2d_like_map); py_bindmethod(type, "apply", array2d_like_apply); + py_bindmethod(type, "zip_with", array2d_like_zip_with); py_bindmethod(type, "copy", array2d_like_copy); py_bindmethod(type, "tolist", array2d_like_tolist); diff --git a/tests/90_array2d.py b/tests/90_array2d.py index 055815de..b683bcb8 100644 --- a/tests/90_array2d.py +++ b/tests/90_array2d.py @@ -24,6 +24,11 @@ assert a.tolist() == [ [(0, 2), (1, 2)], [(0, 3), (1, 3)]] +assert a[0, :].tolist() == [[(0, 0)], [(0, 1)], [(0, 2)], [(0, 3)]] +assert a[1, :].tolist() == [[(1, 0)], [(1, 1)], [(1, 2)], [(1, 3)]] +assert a[:, 0].tolist() == [[(0, 0), (1, 0)]] +assert a[:, -1].tolist() == [[(0, 3), (1, 3)]] + # test is_valid assert a.is_valid(0, 0) and a.is_valid(vec2i(0, 0)) assert a.is_valid(1, 3) and a.is_valid(vec2i(1, 3)) @@ -236,6 +241,12 @@ assert cnt == 1 vis, cnt = a.get_connected_components(0, 'Moore') assert cnt == 2 +# test zip_with +a = array2d[int].fromlist([[1, 2], [3, 4]]) +b = array2d[int].fromlist([[5, 6], [7, 8]]) +c = a.zip_with(b, lambda x, y: x + y) +assert c.tolist() == [[6, 8], [10, 12]] + # stackoverflow bug due to recursive mark-and-sweep # class Cell: # neighbors: list['Cell']