improvements to fam_init for generators that raise

Author: flexatone

arraymap.c

@@ -1189,7 +1189,11 @@ lookup(FAMObject *self, PyObject *key) {
 
 // Insert a key_pos, hash pair into the table. Assumes table already has appropriate size. When inserting a new itme, `hash` is -1, forcing a fresh hash to be computed here. Return 0 on success, -1 on error.
 static int
-insert_obj(FAMObject *self, PyObject *key, Py_ssize_t keys_pos, Py_hash_t hash)
+insert_obj(
+        FAMObject *self,
+        PyObject *key,
+        Py_ssize_t keys_pos,
+        Py_hash_t hash)
 {
     if (hash == -1) {
         hash = PyObject_Hash(key);
@@ -1819,6 +1823,9 @@ fam_init(PyObject *self, PyObject *args, PyObject *kwargs)
 
     if (!keys) {
         keys = PyList_New(0);
+        if (!keys) {
+            return -1;
+        }
     }
     else if (PyObject_TypeCheck(keys, &FAMType)) {
         // Use `keys` as old, `self` as new, and fill from old to new. This returns the same error codes as this function.
@@ -1831,6 +1838,8 @@ fam_init(PyObject *self, PyObject *args, PyObject *kwargs)
             return -1;
         }
         int array_t = PyArray_TYPE(a);
+        keys_size = PyArray_SIZE(a);
+
         if (cls != &AMType &&
                 (PyTypeNum_ISINTEGER(array_t) // signed and unsigned
                 || PyTypeNum_ISFLOAT(array_t)
@@ -1843,25 +1852,26 @@ fam_init(PyObject *self, PyObject *args, PyObject *kwargs)
             keys_array_type = at_to_kat(array_t);
             Py_INCREF(keys);
         }
-        else { // if an AutoMap or an array that we do not custom-hash, we create a list
+        else { // if an AutoMap or an array that we do not handle, create a list
             if (array_t == NPY_DATETIME || array_t == NPY_TIMEDELTA){
                 keys = PySequence_List(keys); // force scalars
             }
             else {
                 keys = PyArray_ToList(a); // converts to objs
             }
+            if (!keys) {
+                return -1;
+            }
         }
-        keys_size = PyArray_SIZE(a);
     }
     else { // assume an arbitrary iterable
         keys = PySequence_List(keys);
+        if (!keys) {
+            return -1;
+        }
         keys_size = PyList_GET_SIZE(keys);
     }
 
-    if (!keys) {
-        return -1;
-    }
-
     fam->keys = keys;
     fam->keys_array_type = keys_array_type;
     fam->keys_size = keys_size;
@@ -1870,7 +1880,6 @@ fam_init(PyObject *self, PyObject *args, PyObject *kwargs)
 
     // NOTE: on itialization, grow_table() does not use keys
     if (grow_table(fam, keys_size)) {
-        // assume `fam->keys` will be decrefed by the caller
         return -1;
     }
     Py_ssize_t i = 0;

test/test_unit.py

@@ -8,8 +8,10 @@
 from arraymap import FrozenAutoMap
 from arraymap import NonUniqueError
 
+
 def test_version():
     import arraymap
+
     assert isinstance(arraymap.__version__, str)
 
 
@@ -41,6 +43,20 @@ def test_fam_contains():
 # ------------------------------------------------------------------------------
 
 
+def test_fam_constructor_a():
+    with pytest.raises(ZeroDivisionError):
+        fam = FrozenAutoMap((x / 0 for x in range(3)))
+
+
+def test_fam_constructor_b():
+    fam1 = FrozenAutoMap(range(3))
+    fam2 = FrozenAutoMap(fam1)
+    assert list(fam2), [0, 1, 2]
+
+
+# ------------------------------------------------------------------------------
+
+
 def test_fam_constructor_array_int_a1():
     a1 = np.array((10, 20, 30), dtype=np.int64)
     with pytest.raises(TypeError):