diff --git a/include/typings/array2d.pyi b/include/typings/array2d.pyi index 5e588256..7a5bf717 100644 --- a/include/typings/array2d.pyi +++ b/include/typings/array2d.pyi @@ -28,6 +28,12 @@ class array2d_like[T]: If the position is out of bounds, returns the default value. """ + def index(self, value: T) -> vec2i: + """Get the position of the first occurrence of the given value. + + Raises `ValueError` if the value is not found. + """ + def render(self) -> str: ... def all(self: array2d_like[bool]) -> bool: ... diff --git a/src/modules/array2d.c b/src/modules/array2d.c index 3df9d468..b1c01110 100644 --- a/src/modules/array2d.c +++ b/src/modules/array2d.c @@ -102,6 +102,24 @@ static bool array2d_like_get(int argc, py_Ref argv) { return true; } +static bool array2d_like_index(int argc, py_Ref argv) { + PY_CHECK_ARGC(2); + c11_array2d_like* self = py_touserdata(argv); + py_Ref value = py_arg(1); + for(int j = 0; j < self->n_rows; j++) { + for(int i = 0; i < self->n_cols; i++) { + py_Ref item = self->f_get(self, i, j); + int code = py_equal(item, value); + if(code == -1) return false; + if(code == 1) { + py_newvec2i(py_retval(), (c11_vec2i){{i, j}}); + return true; + } + } + } + return ValueError("value not found"); +} + static bool array2d_like_render(int argc, py_Ref argv) { PY_CHECK_ARGC(1); c11_sbuf buf; @@ -730,6 +748,7 @@ static void register_array2d_like(py_Ref mod) { py_bindmethod(type, "is_valid", array2d_like_is_valid); py_bindmethod(type, "get", array2d_like_get); + py_bindmethod(type, "index", array2d_like_index); py_bindmethod(type, "render", array2d_like_render); diff --git a/tests/90_array2d.py b/tests/90_array2d.py index 08cf864a..8c78f6fe 100644 --- a/tests/90_array2d.py +++ b/tests/90_array2d.py @@ -12,7 +12,7 @@ except ValueError: pass # test callable constructor -a = array2d[int](2, 4, lambda pos: (pos.x, pos.y)) +a = array2d[tuple[int, int]](2, 4, lambda pos: (pos.x, pos.y)) assert a.width == a.n_cols == 2 assert a.height == a.n_rows == 4 @@ -43,6 +43,10 @@ assert a.get(1, 3) == (1, 3) assert a.get(2, 0) is None assert a.get(0, 4, 'S') == 'S' +# test index +assert a.index((0, 0)) == vec2i(0, 0) +assert a.index((1, 3)) == vec2i(1, 3) + # test __getitem__ assert a[0, 0] == (0, 0) assert a[1, 3] == (1, 3)