blueloveTH há 1 ano atrás
pai
commit
09dc57f206
4 ficheiros alterados com 19 adições e 4 exclusões
  1. 11 4
      src/modules/array2d.c
  2. 2 0
      src/public/internal.c
  3. 1 0
      tests/90_array2d.py
  4. 5 0
      tests/90_chunked_array2d.py

+ 11 - 4
src/modules/array2d.c

@@ -991,12 +991,19 @@ static bool chunked_array2d__delitem__(int argc, py_Ref argv) {
 static bool chunked_array2d__iter__(int argc, py_Ref argv) {
     PY_CHECK_ARGC(1);
     c11_chunked_array2d* self = py_touserdata(argv);
-    py_newtuple(py_retval(), self->chunks.length);
+    py_newtuple(py_pushtmp(), self->chunks.length);
     for(int i = 0; i < self->chunks.length; i++) {
-        py_TValue* data = c11__getitem(c11_chunked_array2d_chunks_KV, &self->chunks, i).value;
-        py_tuple_setitem(py_retval(), i, &data[0]);
+        py_Ref slot = py_tuple_getitem(py_peek(-1), i);
+        c11_chunked_array2d_chunks_KV* kv =
+            c11__at(c11_chunked_array2d_chunks_KV, &self->chunks, i);
+        py_newtuple(slot, 2);
+        py_newvec2i(py_tuple_getitem(slot, 0), kv->key);
+        py_tuple_setitem(slot, 1, &kv->value[0]);
     }
-    return py_iter(py_retval());
+    bool ok = py_iter(py_peek(-1));
+    if(!ok) return false;
+    py_pop();
+    return true;
 }
 
 static bool chunked_array2d__clear(int argc, py_Ref argv) {

+ 2 - 0
src/public/internal.c

@@ -133,6 +133,8 @@ bool py_call(py_Ref f, int argc, py_Ref argv) {
 #ifndef NDEBUG
 bool py_callcfunc(py_CFunction f, int argc, py_Ref argv) {
     py_StackRef p0 = py_peek(0);
+    // NOTE: sometimes users are using `py_retval()` to pass `argv`
+    // It will be reset to `nil` and cause an exception
     py_newnil(py_retval());
     bool ok = f(argc, argv);
     if(!ok) {

+ 1 - 0
tests/90_array2d.py

@@ -16,6 +16,7 @@ a = array2d[int](2, 4, lambda pos: (pos.x, pos.y))
 
 assert a.width == a.n_cols == 2
 assert a.height == a.n_rows == 4
+assert a.shape == vec2i(2, 4)
 assert a.numel == 8
 assert a.tolist() == [
     [(0, 0), (1, 0)],

+ 5 - 0
tests/90_chunked_array2d.py

@@ -3,6 +3,11 @@ from linalg import vec2i
 
 a = chunked_array2d(4, default=0)
 
+print(iter(a))
+print(list(a))
+
 a[vec2i.ONE] = 1
 
 print(a.view().render())
+print(list(a))
+