diff --git a/mypyc/lib-rt/vecs/librt_vecs.c b/mypyc/lib-rt/vecs/librt_vecs.c index 3696e696d5c3e..13d540a1dc451 100644 --- a/mypyc/lib-rt/vecs/librt_vecs.c +++ b/mypyc/lib-rt/vecs/librt_vecs.c @@ -950,36 +950,52 @@ librt_vecs_module_exec(PyObject *m) return -1; if (PyType_Ready(&VecTBufType) < 0) return -1; + if (PyType_Ready(&VecTIterType) < 0) + return -1; if (PyType_Ready(&VecNestedType) < 0) return -1; if (PyType_Ready(&VecNestedBufType) < 0) return -1; + if (PyType_Ready(&VecNestedIterType) < 0) + return -1; if (PyType_Ready(&VecI64Type) < 0) return -1; if (PyType_Ready(&VecI64BufType) < 0) return -1; + if (PyType_Ready(&VecI64IterType) < 0) + return -1; if (PyType_Ready(&VecI32Type) < 0) return -1; if (PyType_Ready(&VecI32BufType) < 0) return -1; + if (PyType_Ready(&VecI32IterType) < 0) + return -1; if (PyType_Ready(&VecI16Type) < 0) return -1; if (PyType_Ready(&VecI16BufType) < 0) return -1; + if (PyType_Ready(&VecI16IterType) < 0) + return -1; if (PyType_Ready(&VecU8Type) < 0) return -1; if (PyType_Ready(&VecU8BufType) < 0) return -1; + if (PyType_Ready(&VecU8IterType) < 0) + return -1; if (PyType_Ready(&VecFloatType) < 0) return -1; if (PyType_Ready(&VecFloatBufType) < 0) return -1; + if (PyType_Ready(&VecFloatIterType) < 0) + return -1; if (PyType_Ready(&VecBoolType) < 0) return -1; if (PyType_Ready(&VecBoolBufType) < 0) return -1; + if (PyType_Ready(&VecBoolIterType) < 0) + return -1; Py_INCREF(&VecType); if (PyModule_AddObject(m, "vec", (PyObject *)&VecType) < 0) { diff --git a/mypyc/lib-rt/vecs/librt_vecs.h b/mypyc/lib-rt/vecs/librt_vecs.h index 80dc977c19146..e9b93157c7982 100644 --- a/mypyc/lib-rt/vecs/librt_vecs.h +++ b/mypyc/lib-rt/vecs/librt_vecs.h @@ -493,6 +493,16 @@ extern PyTypeObject VecBoolType; extern PyTypeObject VecTType; extern PyTypeObject VecNestedType; +// Iterator type objects for vec iteration +extern PyTypeObject VecI64IterType; +extern PyTypeObject VecI32IterType; +extern PyTypeObject VecI16IterType; +extern PyTypeObject VecU8IterType; +extern PyTypeObject VecFloatIterType; +extern PyTypeObject VecBoolIterType; +extern PyTypeObject VecTIterType; +extern PyTypeObject VecNestedIterType; + // Type objects corresponding to the 'i64', 'i32', 'i16, and 'u8' types extern PyTypeObject *LibRTVecs_I64TypeObj; extern PyTypeObject *LibRTVecs_I32TypeObj; diff --git a/mypyc/lib-rt/vecs/vec_nested.c b/mypyc/lib-rt/vecs/vec_nested.c index 0034d9e9bfff5..459a40735744a 100644 --- a/mypyc/lib-rt/vecs/vec_nested.c +++ b/mypyc/lib-rt/vecs/vec_nested.c @@ -420,6 +420,89 @@ static PyMethodDef vec_methods[] = { {NULL, NULL, 0, NULL}, /* Sentinel */ }; +// Iterator type for nested vecs + +typedef struct { + PyObject_HEAD + VecNested vec; // Unboxed vec (keeps buffer alive via buf reference) + Py_ssize_t index; // Current iteration index +} VecNestedIterObject; + +PyTypeObject VecNestedIterType; + +static PyObject *VecNested_iter(PyObject *self) { + VecNestedIterObject *it = PyObject_GC_New(VecNestedIterObject, &VecNestedIterType); + if (it == NULL) + return NULL; + it->vec = ((VecNestedObject *)self)->vec; + Py_INCREF(it->vec.buf); + it->index = 0; + PyObject_GC_Track(it); + return (PyObject *)it; +} + +static int +VecNestedIter_traverse(VecNestedIterObject *self, visitproc visit, void *arg) +{ + Py_VISIT(self->vec.buf); + return 0; +} + +static int +VecNestedIter_clear(VecNestedIterObject *self) +{ + Py_CLEAR(self->vec.buf); + return 0; +} + +static void VecNestedIter_dealloc(VecNestedIterObject *self) { + PyObject_GC_UnTrack(self); + Py_XDECREF(self->vec.buf); + PyObject_GC_Del(self); +} + +static PyObject *VecNestedIter_next(VecNestedIterObject *self) { + if (self->vec.buf == NULL) + return NULL; + if (self->index < self->vec.len) { + PyObject *item = box_vec_item_by_index(self->vec, self->index); + if (item == NULL) + return NULL; + self->index++; + return item; + } + Py_CLEAR(self->vec.buf); + return NULL; // StopIteration +} + +static PyObject *VecNestedIter_len(VecNestedIterObject *self, PyObject *Py_UNUSED(ignored)) { + if (self->vec.buf == NULL) + return PyLong_FromSsize_t(0); + Py_ssize_t remaining = self->vec.len - self->index; + if (remaining < 0) + remaining = 0; + return PyLong_FromSsize_t(remaining); +} + +static PyMethodDef VecNestedIter_methods[] = { + {"__length_hint__", (PyCFunction)VecNestedIter_len, METH_NOARGS, NULL}, + {NULL, NULL, 0, NULL}, +}; + +PyTypeObject VecNestedIterType = { + PyVarObject_HEAD_INIT(NULL, 0) + .tp_name = "vec_nested_iterator", + .tp_basicsize = sizeof(VecNestedIterObject), + .tp_itemsize = 0, + .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, + .tp_traverse = (traverseproc)VecNestedIter_traverse, + .tp_clear = (inquiry)VecNestedIter_clear, + .tp_dealloc = (destructor)VecNestedIter_dealloc, + .tp_iter = PyObject_SelfIter, + .tp_iternext = (iternextfunc)VecNestedIter_next, + .tp_methods = VecNestedIter_methods, +}; + PyTypeObject VecNestedBufType = { PyVarObject_HEAD_INIT(NULL, 0) .tp_name = "vecbuf", @@ -447,6 +530,7 @@ PyTypeObject VecNestedType = { .tp_dealloc = (destructor)VecNested_dealloc, //.tp_free = PyObject_GC_Del, .tp_repr = (reprfunc)vec_repr, + .tp_iter = VecNested_iter, .tp_as_sequence = &VecNestedSequence, .tp_as_mapping = &VecNestedMapping, .tp_richcompare = vec_richcompare, diff --git a/mypyc/lib-rt/vecs/vec_t.c b/mypyc/lib-rt/vecs/vec_t.c index 8a03ea1a2e644..79938cb9aade3 100644 --- a/mypyc/lib-rt/vecs/vec_t.c +++ b/mypyc/lib-rt/vecs/vec_t.c @@ -416,6 +416,88 @@ static PyMethodDef vec_methods[] = { {NULL, NULL, 0, NULL}, /* Sentinel */ }; +// Iterator type for vec[T] (reference types) + +typedef struct { + PyObject_HEAD + VecT vec; // Unboxed vec (keeps buffer alive via buf reference) + Py_ssize_t index; // Current iteration index +} VecTIterObject; + +PyTypeObject VecTIterType; + +static PyObject *VecT_iter(PyObject *self) { + VecTIterObject *it = PyObject_GC_New(VecTIterObject, &VecTIterType); + if (it == NULL) + return NULL; + it->vec = ((VecTObject *)self)->vec; + Py_INCREF(it->vec.buf); + it->index = 0; + PyObject_GC_Track(it); + return (PyObject *)it; +} + +static int +VecTIter_traverse(VecTIterObject *self, visitproc visit, void *arg) +{ + Py_VISIT(self->vec.buf); + return 0; +} + +static int +VecTIter_clear(VecTIterObject *self) +{ + Py_CLEAR(self->vec.buf); + return 0; +} + +static void VecTIter_dealloc(VecTIterObject *self) { + PyObject_GC_UnTrack(self); + Py_XDECREF(self->vec.buf); + PyObject_GC_Del(self); +} + +static PyObject *VecTIter_next(VecTIterObject *self) { + if (self->vec.buf == NULL) + return NULL; + if (self->index < self->vec.len) { + PyObject *item = self->vec.buf->items[self->index]; + self->index++; + Py_INCREF(item); + return item; + } + Py_CLEAR(self->vec.buf); + return NULL; // StopIteration +} + +static PyObject *VecTIter_len(VecTIterObject *self, PyObject *Py_UNUSED(ignored)) { + if (self->vec.buf == NULL) + return PyLong_FromSsize_t(0); + Py_ssize_t remaining = self->vec.len - self->index; + if (remaining < 0) + remaining = 0; + return PyLong_FromSsize_t(remaining); +} + +static PyMethodDef VecTIter_methods[] = { + {"__length_hint__", (PyCFunction)VecTIter_len, METH_NOARGS, NULL}, + {NULL, NULL, 0, NULL}, +}; + +PyTypeObject VecTIterType = { + PyVarObject_HEAD_INIT(NULL, 0) + .tp_name = "vec_iterator", + .tp_basicsize = sizeof(VecTIterObject), + .tp_itemsize = 0, + .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, + .tp_traverse = (traverseproc)VecTIter_traverse, + .tp_clear = (inquiry)VecTIter_clear, + .tp_dealloc = (destructor)VecTIter_dealloc, + .tp_iter = PyObject_SelfIter, + .tp_iternext = (iternextfunc)VecTIter_next, + .tp_methods = VecTIter_methods, +}; + PyTypeObject VecTBufType = { PyVarObject_HEAD_INIT(NULL, 0) .tp_name = "vecbuf", @@ -443,6 +525,7 @@ PyTypeObject VecTType = { .tp_dealloc = (destructor)VecT_dealloc, //.tp_free = PyObject_GC_Del, .tp_repr = (reprfunc)vec_repr, + .tp_iter = VecT_iter, .tp_as_sequence = &VecTSequence, .tp_as_mapping = &VecTMapping, .tp_richcompare = vec_richcompare, diff --git a/mypyc/lib-rt/vecs/vec_template.c b/mypyc/lib-rt/vecs/vec_template.c index dc1e9df3609fc..7881bb496a88c 100644 --- a/mypyc/lib-rt/vecs/vec_template.c +++ b/mypyc/lib-rt/vecs/vec_template.c @@ -354,6 +354,71 @@ static PyMethodDef vec_methods[] = { {NULL, NULL, 0, NULL}, /* Sentinel */ }; +// Iterator type for specialized vec types + +typedef struct { + PyObject_HEAD + VEC vec; // Unboxed vec (keeps buffer alive via buf reference) + Py_ssize_t index; // Current iteration index +} NAME(IterObject); + +PyTypeObject NAME(IterType); + +static PyObject *vec_iter(PyObject *self) { + NAME(IterObject) *it = PyObject_New(NAME(IterObject), &NAME(IterType)); + if (it == NULL) + return NULL; + it->vec = ((VEC_OBJECT *)self)->vec; + Py_XINCREF(it->vec.buf); + it->index = 0; + return (PyObject *)it; +} + +static void vec_iter_dealloc(NAME(IterObject) *self) { + Py_XDECREF(self->vec.buf); + PyObject_Del(self); +} + +static PyObject *vec_iter_next(NAME(IterObject) *self) { + if (self->vec.buf == NULL) + return NULL; + if (self->index < self->vec.len) { + PyObject *item = BOX_ITEM(self->vec.buf->items[self->index]); + if (item == NULL) + return NULL; + self->index++; + return item; + } + Py_CLEAR(self->vec.buf); + return NULL; // StopIteration +} + +static PyObject *vec_iter_len(NAME(IterObject) *self, PyObject *Py_UNUSED(ignored)) { + if (self->vec.buf == NULL) + return PyLong_FromSsize_t(0); + Py_ssize_t remaining = self->vec.len - self->index; + if (remaining < 0) + remaining = 0; + return PyLong_FromSsize_t(remaining); +} + +static PyMethodDef vec_iter_methods[] = { + {"__length_hint__", (PyCFunction)vec_iter_len, METH_NOARGS, NULL}, + {NULL, NULL, 0, NULL}, +}; + +PyTypeObject NAME(IterType) = { + PyVarObject_HEAD_INIT(NULL, 0) + .tp_name = "vec_iterator[" ITEM_TYPE_STR "]", + .tp_basicsize = sizeof(NAME(IterObject)), + .tp_itemsize = 0, + .tp_flags = Py_TPFLAGS_DEFAULT, + .tp_dealloc = (destructor)vec_iter_dealloc, + .tp_iter = PyObject_SelfIter, + .tp_iternext = (iternextfunc)vec_iter_next, + .tp_methods = vec_iter_methods, +}; + PyTypeObject BUF_TYPE = { PyVarObject_HEAD_INIT(NULL, 0) .tp_name = "vecbuf[" ITEM_TYPE_STR "]", @@ -377,6 +442,7 @@ PyTypeObject VEC_TYPE = { //.tp_free = PyObject_Del, .tp_dealloc = (destructor)vec_dealloc, .tp_repr = (reprfunc)vec_repr, + .tp_iter = vec_iter, .tp_as_sequence = &vec_sequence_methods, .tp_as_mapping = &vec_mapping_methods, .tp_richcompare = vec_richcompare, diff --git a/mypyc/test-data/run-vecs-i64-interp.test b/mypyc/test-data/run-vecs-i64-interp.test index 52d816881b66b..166d6818a724c 100644 --- a/mypyc/test-data/run-vecs-i64-interp.test +++ b/mypyc/test-data/run-vecs-i64-interp.test @@ -231,6 +231,45 @@ def test_iteration_nested() -> None: (5, 2), (5, 5), (5, 7), (7, 2), (7, 5), (7, 7)] +def test_iterator_protocol() -> None: + # Test iter() and next() functions explicitly + v = vec[i64]([1, 2, 3]) + it = iter(v) + assert next(it) == 1 + assert next(it) == 2 + assert next(it) == 3 + with assertRaises(StopIteration): + next(it) + + # Test multiple iterators on same vec + v = vec[i64]([10, 20]) + it1 = iter(v) + it2 = iter(v) + assert next(it1) == 10 + assert next(it2) == 10 + assert next(it1) == 20 + assert next(it2) == 20 + + # Test empty vec iterator + it = iter(vec[i64]()) + with assertRaises(StopIteration): + next(it) + +def test_iterator_length_hint() -> None: + v = vec[i64]([1, 2, 3, 4, 5]) + it = iter(v) + op: Any = getattr(sys, "modules")["operator"] + length_hint = op.length_hint + assert length_hint(it) == 5 + next(it) + assert length_hint(it) == 4 + next(it) + next(it) + assert length_hint(it) == 2 + next(it) + next(it) + assert length_hint(it) == 0 + def test_slicing() -> None: v = vec[i64](range(5)) assert v[1:4] == vec[i64]([1, 2, 3]) diff --git a/mypyc/test-data/run-vecs-i64.test b/mypyc/test-data/run-vecs-i64.test index 1c19bfb11e767..0c1fb2f83e548 100644 --- a/mypyc/test-data/run-vecs-i64.test +++ b/mypyc/test-data/run-vecs-i64.test @@ -427,6 +427,29 @@ def test_pop_index() -> None: assert n == 15 assert v == vec[i64]() +def test_iter_and_next() -> None: + v = vec[i64]([1, 2, 3]) + it = iter(v) + assert next(it) == 1 + assert next(it) == 2 + assert next(it) == 3 + with assertRaises(StopIteration): + next(it) + + # Multiple iterators on same vec + v = vec[i64]([10, 20]) + it1 = iter(v) + it2 = iter(v) + assert next(it1) == 10 + assert next(it2) == 10 + assert next(it1) == 20 + assert next(it2) == 20 + + # Empty vec iterator + it = iter(vec[i64]()) + with assertRaises(StopIteration): + next(it) + def f(x: vec[i64]) -> None: pass diff --git a/mypyc/test-data/run-vecs-misc-interp.test b/mypyc/test-data/run-vecs-misc-interp.test index db2f79d926608..07ba78c3f4a29 100644 --- a/mypyc/test-data/run-vecs-misc-interp.test +++ b/mypyc/test-data/run-vecs-misc-interp.test @@ -373,6 +373,86 @@ def test_bool_pop() -> None: assert n is True assert v == vec[bool]([True, False]) +# Iteration tests for all item types + +def test_iteration() -> None: + # Test iteration for each primitive type + for x in vec[float](): + assert False + a = [] + for x in vec[float]([1.5, 2.5, 3.5]): + a.append(x) + assert a == [1.5, 2.5, 3.5] + + for x in vec[u8](): + assert False + a = [] + for x in vec[u8]([1, 2, 255]): + a.append(x) + assert a == [1, 2, 255] + + for x in vec[i16](): + assert False + a = [] + for x in vec[i16]([-1000, 0, 1000]): + a.append(x) + assert a == [-1000, 0, 1000] + + for x in vec[i32](): + assert False + a = [] + for x in vec[i32]([-100000, 0, 100000]): + a.append(x) + assert a == [-100000, 0, 100000] + + for x in vec[bool](): + assert False + a = [] + for x in vec[bool]([True, False, True]): + a.append(x) + assert a == [True, False, True] + +def test_iterator_protocol() -> None: + # Test iter() and next() for float + v_float = vec[float]([1.5, 2.5]) + it = iter(v_float) + assert next(it) == 1.5 + assert next(it) == 2.5 + with assertRaises(StopIteration): + next(it) + + # Test iter() and next() for u8 + v_u8 = vec[u8]([10, 20]) + it = iter(v_u8) + assert next(it) == 10 + assert next(it) == 20 + with assertRaises(StopIteration): + next(it) + + # Test iter() and next() for i16 + v_i16 = vec[i16]([-100, 100]) + it = iter(v_i16) + assert next(it) == -100 + assert next(it) == 100 + with assertRaises(StopIteration): + next(it) + + # Test iter() and next() for i32 + v_i32 = vec[i32]([-1000, 1000]) + it = iter(v_i32) + assert next(it) == -1000 + assert next(it) == 1000 + with assertRaises(StopIteration): + next(it) + + # Test bool specifically since values are True/False singletons + v_bool = vec[bool]([True, False]) + it = iter(v_bool) + assert next(it) is True + assert next(it) is False + with assertRaises(StopIteration): + next(it) + [case testLibrtVecsFeaturesNotAvailableInNonExperimentalBuild_librt] # This also ensures librt.vecs can be built without experimental features import librt.vecs diff --git a/mypyc/test-data/run-vecs-nested-interp.test b/mypyc/test-data/run-vecs-nested-interp.test index a8ed7b87c426c..bd621e99498d1 100644 --- a/mypyc/test-data/run-vecs-nested-interp.test +++ b/mypyc/test-data/run-vecs-nested-interp.test @@ -321,3 +321,92 @@ def test_pop_index() -> None: v, item = pop(v, 0) assert item == vec[str](['15']) assert v == vec[vec[str]]() + +def test_iteration() -> None: + # Test empty nested vec iteration + for x in vec[vec[str]](): + assert False + + # Test single item iteration + a = [] + for x in vec[vec[str]]([vec[str](['a'])]): + a.append(x) + assert a == [vec[str](['a'])] + + # Test multiple items iteration + a = [] + for x in vec[vec[str]]([vec[str](['a']), vec[str](['b', 'c']), vec[str]()]): + a.append(x) + assert a == [vec[str](['a']), vec[str](['b', 'c']), vec[str]()] + + # Test iteration with vec[vec[i64]] + a = [] + for x in vec[vec[i64]]([vec[i64]([1, 2]), vec[i64]([3])]): + a.append(x) + assert a == [vec[i64]([1, 2]), vec[i64]([3])] + +def test_iteration_nested_loop() -> None: + # Test nested iteration over same vec + v = vec[vec[str]]([vec[str](['a']), vec[str](['b'])]) + a = [] + for x in v: + for y in v: + a.append((x, y)) + assert a == [ + (vec[str](['a']), vec[str](['a'])), + (vec[str](['a']), vec[str](['b'])), + (vec[str](['b']), vec[str](['a'])), + (vec[str](['b']), vec[str](['b'])), + ] + +def test_iteration_deeply_nested() -> None: + # Test vec[vec[vec[str]]] iteration + vvv = vec[vec[vec[str]]]([ + vec[vec[str]]([vec[str](['a'])]), + vec[vec[str]]([vec[str](['b']), vec[str](['c'])]) + ]) + a = [] + for x in vvv: + a.append(x) + assert len(a) == 2 + assert a[0] == vec[vec[str]]([vec[str](['a'])]) + assert a[1] == vec[vec[str]]([vec[str](['b']), vec[str](['c'])]) + +def test_iterator_protocol() -> None: + # Test iter() and next() functions explicitly + v = vec[vec[str]]([vec[str](['a']), vec[str](['b']), vec[str](['c'])]) + it = iter(v) + assert next(it) == vec[str](['a']) + assert next(it) == vec[str](['b']) + assert next(it) == vec[str](['c']) + with assertRaises(StopIteration): + next(it) + + # Test multiple iterators on same vec + v = vec[vec[i64]]([vec[i64]([1]), vec[i64]([2])]) + it1 = iter(v) + it2 = iter(v) + assert next(it1) == vec[i64]([1]) + assert next(it2) == vec[i64]([1]) + assert next(it1) == vec[i64]([2]) + assert next(it2) == vec[i64]([2]) + + # Test empty vec iterator + it = iter(vec[vec[str]]()) + with assertRaises(StopIteration): + next(it) + +def test_iterator_length_hint() -> None: + v = vec[vec[str]]([vec[str](['a']), vec[str](['b']), vec[str](['c']), vec[str](['d']), vec[str](['e'])]) + it = iter(v) + op: Any = getattr(sys, "modules")["operator"] + length_hint = op.length_hint + assert length_hint(it) == 5 + next(it) + assert length_hint(it) == 4 + next(it) + next(it) + assert length_hint(it) == 2 + next(it) + next(it) + assert length_hint(it) == 0 diff --git a/mypyc/test-data/run-vecs-nested.test b/mypyc/test-data/run-vecs-nested.test index a8c8032121e7f..0cf430d9c06cc 100644 --- a/mypyc/test-data/run-vecs-nested.test +++ b/mypyc/test-data/run-vecs-nested.test @@ -418,6 +418,18 @@ def test_doubly_nested_vec_i64_operations() -> None: # TODO: box/unbox +def test_iter_and_next() -> None: + a = vec[str](["s1" + str()]) + b = vec[str](["s2" + str()]) + c = vec[str](["s3" + str()]) + v = vec[vec[str]]([a, b, c]) + it = iter(v) + assert next(it) == a + assert next(it) == b + assert next(it) == c + with assertRaises(StopIteration): + next(it) + def f(x: vec[vec[str]]) -> None: pass diff --git a/mypyc/test-data/run-vecs-t-interp.test b/mypyc/test-data/run-vecs-t-interp.test index cb527b50f7c24..881530b50f7a6 100644 --- a/mypyc/test-data/run-vecs-t-interp.test +++ b/mypyc/test-data/run-vecs-t-interp.test @@ -324,6 +324,52 @@ def test_iteration() -> None: a.append(x) assert a == ['11', None, '45', '64'] +def test_iterator_protocol() -> None: + # Test iter() and next() functions explicitly + v = vec[str](['a', 'b', 'c']) + it = iter(v) + assert next(it) == 'a' + assert next(it) == 'b' + assert next(it) == 'c' + with assertRaises(StopIteration): + next(it) + + # Test multiple iterators on same vec + v = vec[str](['x', 'y']) + it1 = iter(v) + it2 = iter(v) + assert next(it1) == 'x' + assert next(it2) == 'x' + assert next(it1) == 'y' + assert next(it2) == 'y' + + # Test empty vec iterator + it = iter(vec[str]()) + with assertRaises(StopIteration): + next(it) + + # Test iteration with optional type + v = vec[Optional[str]](['a', None, 'b']) + it = iter(v) + assert next(it) == 'a' + assert next(it) is None + assert next(it) == 'b' + +def test_iterator_length_hint() -> None: + v = vec[str](['a', 'b', 'c', 'd', 'e']) + it = iter(v) + op: Any = getattr(sys, "modules")["operator"] + length_hint = op.length_hint + assert length_hint(it) == 5 + next(it) + assert length_hint(it) == 4 + next(it) + next(it) + assert length_hint(it) == 2 + next(it) + next(it) + assert length_hint(it) == 0 + def test_slicing() -> None: v = vec[str](str(i) for i in range(5)) assert v[1:4] == vec[str](['1', '2', '3']) diff --git a/mypyc/test-data/run-vecs-t.test b/mypyc/test-data/run-vecs-t.test index 153428148518e..2b69e31408201 100644 --- a/mypyc/test-data/run-vecs-t.test +++ b/mypyc/test-data/run-vecs-t.test @@ -352,6 +352,15 @@ def test_pop_optional() -> None: v, s = pop(v, -1) assert s is None +def test_iter_and_next() -> None: + v = vec[str](["s1" + str(), "s2" + str(), "s3" + str()]) + it = iter(v) + assert next(it) == "s1" + assert next(it) == "s2" + assert next(it) == "s3" + with assertRaises(StopIteration): + next(it) + def f(x: vec[str]) -> None: pass