...

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

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;
 }
 

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);
 

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']