diff --git a/README.rst b/README.rst index e4a357fb..cda1545d 100644 --- a/README.rst +++ b/README.rst @@ -35,6 +35,12 @@ ArrayKit requires the following: What is New in ArrayKit ------------------------- +1.12.0 +............ + +Added ``map_object()`` and ``prepare_iter_for_array()``. + + 1.11.0 ............ diff --git a/src/__init__.py b/src/__init__.py index b33c2a26..7ae728be 100644 --- a/src/__init__.py +++ b/src/__init__.py @@ -28,6 +28,8 @@ from ._arraykit import factorize as factorize from ._arraykit import group_ordering as group_ordering from ._arraykit import group_reduce as group_reduce +from ._arraykit import map_object as map_object +from ._arraykit import prepare_iter_for_array as prepare_iter_for_array from ._arraykit import fill_directional as fill_directional from ._arraykit import count_iteration as count_iteration from ._arraykit import first_true_1d as first_true_1d diff --git a/src/__init__.pyi b/src/__init__.pyi index 7fb69338..aaeb3af3 100644 --- a/src/__init__.pyi +++ b/src/__init__.pyi @@ -236,6 +236,10 @@ def group_ordering( def group_reduce( codes: np.ndarray, size: int, values: np.ndarray, op: str ) -> np.ndarray: ... +def map_object(array: np.ndarray, func: tp.Callable[[tp.Any], tp.Any]) -> np.ndarray: ... +def prepare_iter_for_array( + values: tp.Iterable[tp.Any], copy: bool = ... +) -> tp.Tuple[tp.Optional[type], bool, tp.Sequence[tp.Any]]: ... def fill_directional( array: np.ndarray, target: np.ndarray, diff --git a/src/_arraykit.c b/src/_arraykit.c index 9d81d172..0d14b23f 100644 --- a/src/_arraykit.c +++ b/src/_arraykit.c @@ -82,6 +82,14 @@ static PyMethodDef arraykit_methods[] = { (PyCFunction)group_reduce, METH_VARARGS | METH_KEYWORDS, NULL}, + {"map_object", + (PyCFunction)map_object, + METH_VARARGS | METH_KEYWORDS, + NULL}, + {"prepare_iter_for_array", + (PyCFunction)prepare_iter_for_array, + METH_VARARGS | METH_KEYWORDS, + NULL}, {"fill_directional", (PyCFunction)fill_directional, METH_VARARGS | METH_KEYWORDS, diff --git a/src/methods.c b/src/methods.c index 6ab5fc33..862736ac 100644 --- a/src/methods.c +++ b/src/methods.c @@ -1403,6 +1403,294 @@ group_reduce(PyObject *Py_UNUSED(m), PyObject *args, PyObject *kwargs) return out_arr; } +//------------------------------------------------------------------------------ +// int magnitude beyond which a Python int is no longer losslessly coercible to float; +// mirrors static_frame.core.util.INT_MAX_COERCIBLE_TO_FLOAT +#define AK_INT_MAX_COERCIBLE_TO_FLOAT 1000000000000000LL + +// running state of the prepare_iter_for_array dtype inference over a series of values +typedef struct AK_InferState { + int has_str; + int has_non_str; + int has_inexact; + int has_big_int; + int has_tuple; // a sized object was seen (tuple/list/array/...) + int needs_object; // the resolved dtype must be object rather than auto-detected +} AK_InferState; + +// Classify one value per static_frame.core.util.prepare_iter_for_array's rules, updating +// `s`. Once `s->needs_object` is set the caller can stop inspecting. `enum_type` is +// enum.Enum (may be NULL to skip the Enum check). +static inline void +AK_infer_value(PyObject *v, PyObject *enum_type, AK_InferState *s) +{ + PyTypeObject *vt = Py_TYPE(v); + // exact str/bytes (Python or numpy scalar) -> string; subclasses fall through + if (PyUnicode_CheckExact(v) || PyBytes_CheckExact(v) + || PyArray_IsScalar(v, Unicode) || PyArray_IsScalar(v, String)) { + s->has_str = 1; + } + // a sized object (tuple, list, array, SF container, str subclass) -> object + else if ((vt->tp_as_sequence && vt->tp_as_sequence->sq_length) + || (vt->tp_as_mapping && vt->tp_as_mapping->mp_length)) { + s->has_tuple = 1; + s->needs_object = 1; + } + else { + s->has_non_str = 1; + if (vt == &PyFloat_Type || vt == &PyComplex_Type) { + s->has_inexact = 1; + } + else if (vt == &PyLong_Type) { + int overflow = 0; + long long lv = PyLong_AsLongLongAndOverflow(v, &overflow); + if (overflow || llabs(lv) > AK_INT_MAX_COERCIBLE_TO_FLOAT) { + s->has_big_int = 1; + } + } + else if (PyArray_IsScalar(v, Generic)) { + ; // any other numpy scalar: non-str, no inexact/big-int, not an Enum + } + else if (enum_type != NULL && PyObject_IsInstance(v, enum_type) == 1) { + s->needs_object = 1; + } + } + if ((s->has_str && s->has_non_str) || (s->has_big_int && s->has_inexact)) { + s->needs_object = 1; + } +} + +// Import enum.Enum for the Enum inference case; returns a new reference or NULL +// (clearing the error), in which case the Enum check is skipped. +static PyObject * +AK_import_enum(void) +{ + PyObject *enum_type = NULL; + PyObject *enum_mod = PyImport_ImportModule("enum"); + if (enum_mod != NULL) { + enum_type = PyObject_GetAttrString(enum_mod, "Enum"); + Py_DECREF(enum_mod); + } + if (enum_type == NULL) { + PyErr_Clear(); + } + return enum_type; +} + +static char *map_object_kwarg_names[] = { + "array", + "func", + NULL +}; + +// Apply a Python callable to each element of a 1D array (elements boxed as numpy scalars, +// matching NumPy/Series iteration) and return a new 1D array, inferring the result dtype +// with the same rules as static_frame.core.util.prepare_iter_for_array: the result is an +// object array when the applied values mix strings and non-strings, include a sized object +// (tuple/list/array), an Enum, or mix a large Python int with a Python float/complex; +// otherwise NumPy auto-detects the dtype (e.g. str -> ' float64). This fuses +// the per-element apply, the type inspection, and the array build into one C pass. +PyObject * +map_object(PyObject *Py_UNUSED(m), PyObject *args, PyObject *kwargs) +{ + PyArrayObject *array = NULL; + PyObject *func = NULL; + if (!PyArg_ParseTupleAndKeywords(args, kwargs, + "O!O:map_object", + map_object_kwarg_names, + &PyArray_Type, &array, + &func + )) { + return NULL; + } + if (PyArray_NDIM(array) != 1) { + PyErr_SetString(PyExc_ValueError, "array must be 1-dimensional"); + return NULL; + } + if (!PyCallable_Check(func)) { + PyErr_SetString(PyExc_TypeError, "func must be callable"); + return NULL; + } + npy_intp n = PyArray_SIZE(array); + int is_object = PyArray_TYPE(array) == NPY_OBJECT; + + PyObject *values = PyList_New(n); // collected results; owns references + if (values == NULL) { + return NULL; + } + // enum.Enum for the rare Enum-result case; on failure proceed without the check + PyObject *enum_type = AK_import_enum(); + AK_InferState state = {0, 0, 0, 0, 0, 0}; + + // 1D: hoist the base pointer and element stride and walk a running pointer, rather + // than recomputing PyArray_GETPTR1 each iteration. For a contiguous array the stride + // is the itemsize (direct indexing into the flat buffer); a strided slice still works. + char *p = (char*)PyArray_DATA(array); + npy_intp stride = PyArray_STRIDES(array)[0]; + + for (npy_intp i = 0; i < n; i++, p += stride) { + PyObject *elem; + if (is_object) { + elem = *(PyObject**)p; + Py_INCREF(elem); + } + else { + elem = PyArray_ToScalar(p, array); + if (elem == NULL) { + goto fail; + } + } + PyObject *r = PyObject_CallOneArg(func, elem); + Py_DECREF(elem); + if (r == NULL) { + goto fail; + } + PyList_SET_ITEM(values, i, r); // steals reference to r + if (!state.needs_object) { + AK_infer_value(r, enum_type, &state); + } + } + Py_XDECREF(enum_type); + + PyObject *result; + if (state.needs_object) { + // build an object array of the collected values + result = PyArray_FROM_OTF(values, NPY_OBJECT, NPY_ARRAY_C_CONTIGUOUS); + } + else { + // let NumPy auto-detect the dtype from the values (str -> ' f8, ...) + result = PyArray_FromAny(values, NULL, 1, 1, NPY_ARRAY_C_CONTIGUOUS, NULL); + } + Py_DECREF(values); + if (result == NULL) { + return NULL; + } + PyArray_CLEARFLAGS((PyArrayObject*)result, NPY_ARRAY_WRITEABLE); + return result; + +fail: + Py_XDECREF(enum_type); + Py_DECREF(values); + return NULL; +} + +static char *prepare_iter_for_array_kwarg_names[] = { + "values", + "copy", + NULL +}; + +// Infer a dtype specifier for the elements of an iterable, matching +// static_frame.core.util.prepare_iter_for_array: return ``(resolved, has_tuple, values)`` +// where `resolved` is None (let NumPy auto-detect) or the ``object`` type, `has_tuple` +// marks that a sized object was seen, and `values` is a newly materialized list when +// `copy` is true (e.g. a generator/dict/set) else the original iterable. The caller +// decides `copy` (via is_gen_copy_values), keeping the SF-specific type policy out of C. +PyObject * +prepare_iter_for_array(PyObject *Py_UNUSED(m), PyObject *args, PyObject *kwargs) +{ + PyObject *values = NULL; + int copy = 0; + if (!PyArg_ParseTupleAndKeywords(args, kwargs, + "O|p:prepare_iter_for_array", + prepare_iter_for_array_kwarg_names, + &values, + © + )) { + return NULL; + } + PyObject *enum_type = AK_import_enum(); + AK_InferState state = {0, 0, 0, 0, 0, 0}; + PyObject *values_out = NULL; // new list when copy, else a new ref to the original + + if (copy) { + // materialize into a list while inspecting (generator/dict/set input). Pre-size + // the list from a length hint when one is available (sets, dicts, sized + // iterators) to avoid repeated reallocation; fall back to append growth for a + // bare generator (hint 0), and guard against an inexact hint over/under-shooting. + Py_ssize_t hint = PyObject_LengthHint(values, 0); + if (hint < 0) { + goto fail; + } + values_out = PyList_New(hint); + if (values_out == NULL) { + goto fail; + } + PyObject *iter = PyObject_GetIter(values); + if (iter == NULL) { + goto fail; + } + Py_ssize_t i = 0; + PyObject *item; + while ((item = PyIter_Next(iter)) != NULL) { + if (!state.needs_object) { + AK_infer_value(item, enum_type, &state); + } + if (i < hint) { + PyList_SET_ITEM(values_out, i, item); // steals reference + } + else { // hint underestimated the length + int rc = PyList_Append(values_out, item); + Py_DECREF(item); + if (rc != 0) { + Py_DECREF(iter); + goto fail; + } + } + i++; + } + Py_DECREF(iter); + if (PyErr_Occurred()) { + goto fail; + } + if (i < hint) { // hint overestimated: drop the trailing (NULL) slots + Py_SET_SIZE(values_out, i); + } + } + else { + // inspect only; direct indexing for list/tuple, else the iterator protocol + Py_INCREF(values); + values_out = values; + if (PyList_CheckExact(values) || PyTuple_CheckExact(values)) { + Py_ssize_t sz = PySequence_Fast_GET_SIZE(values); + for (Py_ssize_t i = 0; i < sz && !state.needs_object; i++) { + AK_infer_value(PySequence_Fast_GET_ITEM(values, i), enum_type, &state); + } + } + else { + PyObject *iter = PyObject_GetIter(values); + if (iter == NULL) { + goto fail; + } + PyObject *item; + while (!state.needs_object && (item = PyIter_Next(iter)) != NULL) { + AK_infer_value(item, enum_type, &state); + Py_DECREF(item); + } + Py_DECREF(iter); + if (PyErr_Occurred()) { + goto fail; + } + } + } + Py_XDECREF(enum_type); + + PyObject *resolved = state.needs_object + ? (PyObject*)&PyBaseObject_Type // the ``object`` builtin -> object dtype + : Py_None; + PyObject *has_tuple = state.has_tuple ? Py_True : Py_False; + PyObject *result = PyTuple_Pack(3, resolved, has_tuple, values_out); + Py_DECREF(values_out); + return result; + +fail: + Py_XDECREF(enum_type); + Py_XDECREF(values_out); + return NULL; +} + +//------------------------------------------------------------------------------ + // Fill one strided lane in place: walk positions in the fill direction, carrying // the most recent non-target value into each target position (subject to `limit` // consecutive fills per run). `elem_base`/`elem_stride` address elements in bytes; diff --git a/src/methods.h b/src/methods.h index e941b178..984457ef 100644 --- a/src/methods.h +++ b/src/methods.h @@ -75,6 +75,12 @@ group_ordering(PyObject *Py_UNUSED(m), PyObject *args, PyObject *kwargs); PyObject * group_reduce(PyObject *Py_UNUSED(m), PyObject *args, PyObject *kwargs); +PyObject * +map_object(PyObject *Py_UNUSED(m), PyObject *args, PyObject *kwargs); + +PyObject * +prepare_iter_for_array(PyObject *Py_UNUSED(m), PyObject *args, PyObject *kwargs); + PyObject * fill_directional(PyObject *Py_UNUSED(m), PyObject *args, PyObject *kwargs); diff --git a/test/test_map_object.py b/test/test_map_object.py new file mode 100644 index 00000000..68fcdecc --- /dev/null +++ b/test/test_map_object.py @@ -0,0 +1,153 @@ +import unittest +from enum import Enum + +import numpy as np +from arraykit import map_object + + +class Color(Enum): + R = 1 + G = 2 + + +# reference: prepare_iter_for_array's inference + build, as static_frame applies it +_INEXACT = (float, complex, np.inexact) +_BIG = 1_000_000_000_000_000 + + +def _reference(arr, func): + vals = [func(v) for v in arr] + resolved = None + has_str = has_non = has_inx = has_big = False + for v in vals: + vt = v.__class__ + if vt is str or vt is np.str_ or vt is bytes or vt is np.bytes_: + has_str = True + elif hasattr(v, '__len__') or isinstance(v, Enum): + resolved = object + break + else: + has_non = True + if vt in _INEXACT: + has_inx = True + elif vt is int and abs(v) > _BIG: + has_big = True + if (has_str and has_non) or (has_big and has_inx): + resolved = object + break + return np.array(vals) if resolved is None else np.array(vals, dtype=object) + + +class TestUnit(unittest.TestCase): + def _check(self, arr, func): + post = map_object(arr, func) + exp = _reference(arr, func) + self.assertEqual(post.dtype, exp.dtype, (arr.dtype, exp.dtype)) + self.assertTrue(np.array_equal(post, exp)) + self.assertFalse(post.flags.writeable) + return post + + def test_map_object_str_from_float(self) -> None: + post = self._check(np.array([1.5, 2.25, 3.0]), lambda x: str(x)) + self.assertEqual(post.dtype.kind, 'U') + + def test_map_object_str_from_bool(self) -> None: + post = self._check(np.array([True, False, True]), lambda x: str(x)) + self.assertEqual(post.tolist(), ['True', 'False', 'True']) + self.assertEqual(post.dtype.kind, 'U') + + def test_map_object_str_from_int(self) -> None: + self._check(np.array([1, 2, 3], dtype=np.int64), lambda x: str(x)) + + def test_map_object_native_float(self) -> None: + post = self._check(np.array([1.5, 2.5]), lambda x: float(x) * 2) + self.assertEqual(post.dtype, np.dtype(np.float64)) + + def test_map_object_native_int(self) -> None: + # python-int results auto-detect to the platform default int (int32 on Windows) + post = self._check(np.array([1, 2, 3]), lambda x: int(x) + 1) + self.assertEqual(post.dtype, np.dtype(np.int_)) + + def test_map_object_tuple_result(self) -> None: + post = self._check(np.array([1, 2]), lambda x: (int(x), int(x))) + self.assertEqual(post.dtype, np.dtype(object)) + + def test_map_object_list_result(self) -> None: + post = self._check(np.array([1, 2]), lambda x: [int(x)]) + self.assertEqual(post.dtype, np.dtype(object)) + + def test_map_object_mixed_str_nonstr(self) -> None: + post = self._check(np.array([1, 2, 3]), lambda x: str(x) if x > 1 else int(x)) + self.assertEqual(post.dtype, np.dtype(object)) + + def test_map_object_python_float(self) -> None: + self._check(np.array([1, 2, 3]), lambda x: 1.5) + + def test_map_object_bigint_and_inexact(self) -> None: + # a large python int mixed with a python float -> object + post = self._check(np.array([1, 2]), lambda x: 10**18 if x == 1 else 1.5) + self.assertEqual(post.dtype, np.dtype(object)) + + def test_map_object_bigint_only(self) -> None: + # big ints alone (no inexact) do not force object + post = self._check(np.array([1, 2]), lambda x: 10**18) + self.assertNotEqual(post.dtype, np.dtype(object)) + + def test_map_object_enum_result(self) -> None: + post = self._check(np.array([1, 2]), lambda x: Color.R) + self.assertEqual(post.dtype, np.dtype(object)) + + def test_map_object_object_input(self) -> None: + arr = np.array(['a', 'bb', 'ccc'], dtype=object) + post = self._check(arr, lambda x: len(x)) + self.assertEqual(post.tolist(), [1, 2, 3]) + + def test_map_object_receives_numpy_scalar(self) -> None: + # elements are boxed as numpy scalars, matching Series/array iteration + seen = [] + map_object(np.array([1.5, 2.5]), lambda x: seen.append(type(x)) or x) + self.assertTrue(all(t is np.float64 for t in seen)) + + def test_map_object_str_subclass_is_object(self) -> None: + # a str subclass is not an exact str -> sized object -> object array + class S(str): + pass + + post = map_object(np.array([1, 2]), lambda x: S(str(x))) + self.assertEqual(post.dtype, np.dtype(object)) + + def test_map_object_empty(self) -> None: + post = self._check(np.array([], dtype=np.float64), lambda x: str(x)) + self.assertEqual(len(post), 0) + + def test_map_object_strided_non_contiguous(self) -> None: + # a strided slice (non-contiguous) must be walked correctly by the running pointer + base = np.array([1.0, 99.0, 2.0, 99.0, 3.0]) + strided = base[::2] + self.assertFalse(strided.flags['C_CONTIGUOUS']) + post = self._check(strided, lambda x: str(x)) + self.assertEqual(post.tolist(), ['1.0', '2.0', '3.0']) + + def test_map_object_strided_object(self) -> None: + arr = np.array(['a', 'X', 'bb', 'X', 'ccc'], dtype=object)[::2] + post = self._check(arr, lambda x: len(x)) + self.assertEqual(post.tolist(), [1, 2, 3]) + + def test_map_object_propagates_exception(self) -> None: + def bad(x): + raise ValueError('boom') + + with self.assertRaises(ValueError): + map_object(np.array([1, 2]), bad) + + def test_map_object_errors(self) -> None: + with self.assertRaises(ValueError): # 2d + map_object(np.array([[1, 2]]), lambda x: x) + with self.assertRaises(TypeError): # not callable + map_object(np.array([1, 2]), 3) + with self.assertRaises(TypeError): # not an array + map_object([1, 2], lambda x: x) + + +if __name__ == '__main__': + unittest.main() diff --git a/test/test_prepare_iter_for_array.py b/test/test_prepare_iter_for_array.py new file mode 100644 index 00000000..8af2155f --- /dev/null +++ b/test/test_prepare_iter_for_array.py @@ -0,0 +1,190 @@ +import unittest +from enum import Enum + +import numpy as np +from arraykit import prepare_iter_for_array + + +class Color(Enum): + R = 1 + G = 2 + + +class _Hinted: + """An iterable exposing a (possibly inexact) __length_hint__.""" + + def __init__(self, items, hint): + self._items = list(items) + self._hint = hint + + def __iter__(self): + return iter(self._items) + + def __length_hint__(self): + return self._hint + + +# faithful reference: static_frame.core.util.prepare_iter_for_array, given a precomputed +# copy flag (SF's is_gen_copy_values decision lives in the SF wrapper) +_INEXACT = (float, complex, np.inexact) +_BIG = 1_000_000_000_000_000 + + +def _reference(values, copy): + if copy: + vpost = [] + resolved = None + has_tuple = False + has_str = has_non = has_inx = has_big = False + it = iter(values) + for v in it: + if copy: + vpost.append(v) + vt = v.__class__ + if vt is str or vt is np.str_ or vt is bytes or vt is np.bytes_: + has_str = True + elif hasattr(v, '__len__'): + has_tuple = True + resolved = object + break + elif isinstance(v, Enum): + resolved = object + break + else: + has_non = True + if vt in _INEXACT: + has_inx = True + elif vt is int and abs(v) > _BIG: + has_big = True + if (has_str and has_non) or (has_big and has_inx): + resolved = object + break + if copy: + vpost.extend(it) + return resolved, has_tuple, vpost + return resolved, has_tuple, values + + +class TestUnit(unittest.TestCase): + def _check(self, make, copy): + # make() returns a fresh iterable so the two runs are independent + r_ak = prepare_iter_for_array(make(), copy) + r_ref = _reference(make(), copy) + self.assertIs(r_ak[0], r_ref[0]) # None or the object type, by identity + self.assertEqual(r_ak[1], r_ref[1]) # has_tuple + self.assertEqual(list(r_ak[2]), list(r_ref[2])) # values + return r_ak + + def test_list_str(self) -> None: + r = self._check(lambda: ['a', 'b', 'c'], False) + self.assertIsNone(r[0]) + + def test_list_float(self) -> None: + self._check(lambda: [1.0, 2.0, 3.0], False) + + def test_tuple_int(self) -> None: + self._check(lambda: (1, 2, 3), False) + + def test_mixed_str_nonstr(self) -> None: + r = self._check(lambda: [1, 'a', 2.0], False) + self.assertIs(r[0], object) + + def test_sized_object_has_tuple(self) -> None: + r = self._check(lambda: [1, (2, 3)], False) + self.assertIs(r[0], object) + self.assertTrue(r[1]) # has_tuple + + def test_enum(self) -> None: + r = self._check(lambda: [Color.R, Color.G], False) + self.assertIs(r[0], object) + self.assertFalse(r[1]) # not has_tuple + + def test_bigint_and_inexact(self) -> None: + r = self._check(lambda: [10**18, 1.5], False) + self.assertIs(r[0], object) + + def test_bigint_only(self) -> None: + r = self._check(lambda: [10**18, 2], False) + self.assertIsNone(r[0]) + + def test_numpy_float_scalars(self) -> None: + r = self._check(lambda: [np.float64(1.5), np.float64(2.5)], False) + self.assertIsNone(r[0]) + + def test_bytes(self) -> None: + self._check(lambda: [b'x', b'y'], False) + + def test_empty_list(self) -> None: + r = self._check(list, False) + self.assertIsNone(r[0]) + + def test_generator_copy(self) -> None: + # a generator is materialized when copy=True; the returned list is the values + r = self._check(lambda: (str(i) for i in range(4)), True) + self.assertEqual(list(r[2]), ['0', '1', '2', '3']) + self.assertIsInstance(r[2], list) + + def test_generator_copy_mixed(self) -> None: + r = self._check(lambda: (i if i < 2 else str(i) for i in range(4)), True) + self.assertIs(r[0], object) + self.assertEqual(list(r[2]), [0, 1, '2', '3']) + + def test_generator_empty_copy(self) -> None: + self._check(lambda: (x for x in []), True) + + def test_set_copy(self) -> None: + # a set is materialized (order-independent check of contents); it has __len__, + # so the list is pre-sized + r = prepare_iter_for_array({1, 2, 3}, True) + self.assertIsNone(r[0]) + self.assertEqual(sorted(r[2]), [1, 2, 3]) + + def test_dict_copy(self) -> None: + r = prepare_iter_for_array({'a': 1, 'b': 2}, True) + self.assertEqual(sorted(r[2]), ['a', 'b']) + + def test_length_hint_exact(self) -> None: + r = self._check(lambda: _Hinted([1, 2, 3], 3), True) + self.assertEqual(list(r[2]), [1, 2, 3]) + + def test_length_hint_overestimate(self) -> None: + # a hint larger than the actual length -> trailing slots dropped + r = prepare_iter_for_array(_Hinted([1, 2, 3], 10), True) + self.assertEqual(list(r[2]), [1, 2, 3]) + self.assertEqual(len(r[2]), 3) + + def test_length_hint_underestimate(self) -> None: + # a hint smaller than the actual length -> remaining items appended + r = prepare_iter_for_array(_Hinted([1, 2, 3, 4, 5], 1), True) + self.assertEqual(list(r[2]), [1, 2, 3, 4, 5]) + + def test_length_hint_inference_preserved(self) -> None: + # inference still resolves object through the pre-sized path + r = prepare_iter_for_array(_Hinted([1, (2,)], 5), True) + self.assertIs(r[0], object) + self.assertTrue(r[1]) + self.assertEqual(list(r[2]), [1, (2,)]) + + def test_no_copy_returns_original(self) -> None: + src = [1, 2, 3] + r = prepare_iter_for_array(src, False) + self.assertIs(r[2], src) # original object, not a copy + + def test_early_stop_does_not_over_iterate(self) -> None: + # once object is resolved, inspection stops; a later error-raising element in a + # non-copy list is never inspected + r = prepare_iter_for_array([1, (2,), object()], False) + self.assertIs(r[0], object) + self.assertTrue(r[1]) + + def test_generator_copy_propagates_exception(self) -> None: + def gen(): + yield 1 + raise ValueError('boom') + + with self.assertRaises(ValueError): + prepare_iter_for_array(gen(), True) + + +if __name__ == '__main__': + unittest.main()