diff --git a/include/typings/array2d.pyi b/include/typings/array2d.pyi index 4c37a093..ea115aa1 100644 --- a/include/typings/array2d.pyi +++ b/include/typings/array2d.pyi @@ -38,6 +38,7 @@ class array2d(Generic[T]): def fill_(self, value: T) -> None: ... def apply_(self, f: Callable[[T], T]) -> None: ... + def indexed_apply_(self, f: Callable[[int, int, T], T]) -> None: ... def copy_(self, other: 'array2d[T] | list[T]') -> None: ... # algorithms diff --git a/src/array2d.cpp b/src/array2d.cpp index 73dc7843..89e457d1 100644 --- a/src/array2d.cpp +++ b/src/array2d.cpp @@ -229,6 +229,16 @@ struct Array2d{ return vm->None; }); + vm->bind(type, "indexed_apply_(self, f)", [](VM* vm, ArgsView args){ + Array2d& self = PK_OBJ_GET(Array2d, args[0]); + PyObject* f = args[1]; + for(int i = 0; i < self.numel; i++){ + std::div_t res = std::div(i, self.n_cols); + self.data[i] = vm->call(f, VAR(res.rem), VAR(res.quot), self.data[i]); + } + return vm->None; + }); + vm->bind(type, "copy_(self, other)", [](VM* vm, ArgsView args){ Array2d& self = PK_OBJ_GET(Array2d, args[0]); if(is_type(args[1], VM::tp_list)){ diff --git a/tests/83_array2d.py b/tests/83_array2d.py index a35cf3da..2444be30 100644 --- a/tests/83_array2d.py +++ b/tests/83_array2d.py @@ -166,3 +166,9 @@ try: exit(1) except TypeError: pass + +a = array2d(3, 4) +a.indexed_apply_(lambda x, y, val: x+y) +assert a[0, 0] == 0 +assert a[1, 2] == 3 +assert a[2, 0] == 2