Merge pull request #42 from InvestmentSystems/39/array-deepcopy-kwarg

`array_deepcopy` interface to accept kwargs, memo dict optional

Author: flexatone

src/_arraykit.c

@@ -148,22 +148,25 @@ AK_ResolveDTypeIter(PyObject *dtypes)
     return resolved;
 }
 
-// Numpy implementation: https://github.com/numpy/numpy/blob/a14c41264855e44ebd6187d7541b5b8d59bb32cb/numpy/core/src/multiarray/methods.c#L1557
+// Perform a deepcopy on an array, using an optional memo dictionary, and specialized to depend on immutable arrays. Related Numpy implementation: https://github.com/numpy/numpy/blob/a14c41264855e44ebd6187d7541b5b8d59bb32cb/numpy/core/src/multiarray/methods.c#L1557
 PyObject*
 AK_ArrayDeepCopy(PyArrayObject *array, PyObject *memo)
 {
     PyObject *id = PyLong_FromVoidPtr((PyObject*)array);
     if (!id) {
         return NULL;
     }
-    PyObject *found = PyDict_GetItemWithError(memo, id);
-    if (found) { // found will be NULL if not in dict
-        Py_INCREF(found); // got a borrowed ref, increment first
-        Py_DECREF(id);
-        return found;
-    }
-    else if (PyErr_Occurred()) {
-        goto error;
+
+    if (memo) {
+        PyObject *found = PyDict_GetItemWithError(memo, id);
+        if (found) { // found will be NULL if not in dict
+            Py_INCREF(found); // got a borrowed ref, increment first
+            Py_DECREF(id);
+            return found;
+        }
+        else if (PyErr_Occurred()) {
+            goto error;
+        }
     }
 
     // if dtype is object, call deepcopy with memo
@@ -187,14 +190,17 @@ AK_ArrayDeepCopy(PyArrayObject *array, PyObject *memo)
         }
     }
     else {
+        // if not a n object dtype, we will force a copy (even if this is an immutable array) so as to not hold on to any references
         Py_INCREF(dtype); // PyArray_FromArray steals a reference
         array_new = PyArray_FromArray(
                 array,
                 dtype,
                 NPY_ARRAY_ENSURECOPY);
-        if (!array_new || PyDict_SetItem(memo, id, array_new)) {
-            Py_XDECREF(array_new);
-            goto error;
+        if (memo) {
+            if (!array_new || PyDict_SetItem(memo, id, array_new)) {
+                Py_XDECREF(array_new);
+                goto error;
+            }
         }
     }
     // set immutable
@@ -311,20 +317,25 @@ row_1d_filter(PyObject *Py_UNUSED(m), PyObject *a)
 //------------------------------------------------------------------------------
 // array utility
 
-// Specialized array deepcopy that stores immutable arrays in memo dict.
+static char *array_deepcopy_kwarg_names[] = {
+    "array",
+    "memo",
+    NULL
+};
+
+// Specialized array deepcopy that stores immutable arrays in an optional memo dict that can be provided with kwargs.
 static PyObject *
-array_deepcopy(PyObject *Py_UNUSED(m), PyObject *args)
+array_deepcopy(PyObject *Py_UNUSED(m), PyObject *args, PyObject *kwargs)
 {
-    PyObject *array, *memo;
-    if (!PyArg_UnpackTuple(args, "array_deepcopy", 2, 2, &array, &memo)) {
+    PyObject *array;
+    PyObject *memo = NULL;
+    if (!PyArg_ParseTupleAndKeywords(args, kwargs,
+            "O|O!:array_deepcopy", array_deepcopy_kwarg_names,
+            &array,
+            &PyDict_Type, &memo)) {
         return NULL;
     }
     AK_CHECK_NUMPY_ARRAY(array);
-    if (!PyDict_CheckExact(memo)) {
-        PyErr_Format(PyExc_TypeError, "expected a dict (got %s)",
-                Py_TYPE(memo)->tp_name);
-        return NULL;
-    }
     return AK_ArrayDeepCopy((PyArrayObject*)array, memo);
 }
 
@@ -623,7 +634,10 @@ static PyMethodDef arraykit_methods[] =  {
     {"column_2d_filter", column_2d_filter, METH_O, NULL},
     {"column_1d_filter", column_1d_filter, METH_O, NULL},
     {"row_1d_filter", row_1d_filter, METH_O, NULL},
-    {"array_deepcopy", array_deepcopy, METH_VARARGS, NULL},
+    {"array_deepcopy",
+            (PyCFunction)array_deepcopy,
+            METH_VARARGS | METH_KEYWORDS,
+            NULL},
     {"resolve_dtype", resolve_dtype, METH_VARARGS, NULL},
     {"resolve_dtype_iter", resolve_dtype_iter, METH_O, NULL},
     {NULL},

test/test_util.py

@@ -177,7 +177,7 @@ def test_array_deepcopy_a1(self) -> None:
         memo = {}
         a2 = array_deepcopy(a1, memo)
 
-        self.assertNotEqual(id(a1), id(a2))
+        self.assertIsNot(a1, a2)
         self.assertNotEqual(mloc(a1), mloc(a2))
         self.assertFalse(a2.flags.writeable)
         self.assertEqual(a1.dtype, a2.dtype)
@@ -187,7 +187,7 @@ def test_array_deepcopy_a2(self) -> None:
         memo = {}
         a2 = array_deepcopy(a1, memo)
 
-        self.assertNotEqual(id(a1), id(a2))
+        self.assertIsNot(a1, a2)
         self.assertNotEqual(mloc(a1), mloc(a2))
         self.assertIn(id(a1), memo)
         self.assertEqual(memo[id(a1)].tolist(), a2.tolist())
@@ -208,23 +208,43 @@ def test_array_deepcopy_c1(self) -> None:
         a1 = np.array((None, 'foo', True, mutable))
         a2 = array_deepcopy(a1, memo)
 
-        self.assertNotEqual(id(a1), id(a2))
+        self.assertIsNot(a1, a2)
         self.assertNotEqual(mloc(a1), mloc(a2))
-        self.assertNotEqual(id(a1[3]), id(a2[3]))
+        self.assertIsNot(a1[3], a2[3])
         self.assertFalse(a2.flags.writeable)
 
     def test_array_deepcopy_c2(self) -> None:
         memo = {}
         mutable = [np.nan]
         a1 = np.array((None, 'foo', True, mutable))
         a2 = array_deepcopy(a1, memo)
-        self.assertNotEqual(id(a1), id(a2))
+        self.assertIsNot(a1, a2)
         self.assertNotEqual(mloc(a1), mloc(a2))
-        self.assertNotEqual(id(a1[3]), id(a2[3]))
+        self.assertIsNot(a1[3], a2[3])
         self.assertFalse(a2.flags.writeable)
         self.assertIn(id(a1), memo)
 
 
+    def test_array_deepcopy_d(self) -> None:
+        memo = {}
+        mutable = [3, 4, 5]
+        a1 = np.array((None, 'foo', True, mutable))
+        a2 = array_deepcopy(a1, memo=memo)
+        self.assertIsNot(a1, a2)
+        self.assertTrue(id(mutable) in memo)
+
+    def test_array_deepcopy_e(self) -> None:
+        a1 = np.array((3, 4, 5))
+        with self.assertRaises(TypeError):
+            # memo argument must be a dictionary
+            a2 = array_deepcopy(a1, memo=None)
+
+    def test_array_deepcopy_f(self) -> None:
+        a1 = np.array((3, 4, 5))
+        a2 = array_deepcopy(a1)
+        self.assertNotEqual(id(a1), id(a2))
+
+
 if __name__ == '__main__':
     unittest.main()