TST: Add some very basic tests to ensure coverage

3 files changed, 125 insertions, 1 deletions

diff --git a/numpy/core/setup.py b/numpy/core/setup.py
index 8d2a86343..2061eb510 100644
--- a/numpy/core/setup.py
+++ b/numpy/core/setup.py
@@ -701,9 +701,11 @@ def configuration(parent_package='',top_path=None):
     config.add_extension('_multiarray_tests',
                     sources=[join('src', 'multiarray', '_multiarray_tests.c.src'),
                              join('src', 'common', 'mem_overlap.c'),
-                             join('src', 'common', 'npy_argparse.c')],
+                             join('src', 'common', 'npy_argparse.c'),
+                             join('src', 'common', 'npy_hashtable.c')],
                     depends=[join('src', 'common', 'mem_overlap.h'),
                              join('src', 'common', 'npy_argparse.h'),
+                             join('src', 'common', 'npy_hashtable.h'),
                              join('src', 'common', 'npy_extint128.h')],
                     libraries=['npymath'])
 
diff --git a/numpy/core/src/multiarray/_multiarray_tests.c.src b/numpy/core/src/multiarray/_multiarray_tests.c.src
index 79140bdb7..4b9a4f9dd 100644
--- a/numpy/core/src/multiarray/_multiarray_tests.c.src
+++ b/numpy/core/src/multiarray/_multiarray_tests.c.src
@@ -1,4 +1,6 @@
 /* -*-c-*- */
+#define PY_SSIZE_T_CLEAN
+
 #define NPY_NO_DEPRECATED_API NPY_API_VERSION
 #include <Python.h>
 #define _NPY_NO_DEPRECATIONS /* for NPY_CHAR */
@@ -11,6 +13,7 @@
 #include "mem_overlap.h"
 #include "npy_extint128.h"
 #include "array_method.h"
+#include "npy_hashtable.h"
 
 #if defined(MS_WIN32) || defined(__CYGWIN__)
 #define EXPORT(x) __declspec(dllexport) x
@@ -1116,6 +1119,92 @@ get_all_cast_information(PyObject *NPY_UNUSED(mod), PyObject *NPY_UNUSED(args))
 
 
 /*
+ * Helper to test the identity cache, takes a list of values and adds
+ * all to the cache except the last key/value pair.  The last value is
+ * ignored, instead the last key is looked up.
+ * None is returned, if the key is not found.
+ * If `replace` is True, duplicate entries are ignored when adding to the
+ * hashtable.
+ */
+static PyObject *
+identityhash_tester(PyObject *NPY_UNUSED(mod),
+        PyObject *const *args, Py_ssize_t len_args, PyObject *kwnames)
+{
+    NPY_PREPARE_ARGPARSER;
+
+    int key_len;
+    int replace;
+    PyObject *replace_obj = Py_False;
+    PyObject *sequence;
+    PyObject *result = NULL;
+
+    if (npy_parse_arguments("identityhash_tester", args, len_args, kwnames,
+            "key_len", &PyArray_PythonPyIntFromInt, &key_len,
+            "sequence", NULL, &sequence,
+            "|replace", NULL, &replace_obj,
+            NULL, NULL, NULL) < 0) {
+        return NULL;
+    }
+    replace = PyObject_IsTrue(replace_obj);
+    if (error_converting(replace)) {
+        return NULL;
+    }
+
+    if (key_len < 1 || key_len >= NPY_MAXARGS) {
+        PyErr_SetString(PyExc_ValueError, "must have 1 to max-args keys.");
+        return NULL;
+    }
+    PyArrayIdentityHash *tb = PyArrayIdentityHash_New(key_len);
+    if (tb == NULL) {
+        return NULL;
+    }
+
+    /* Replace the sequence with a guaranteed fast-sequence */
+    sequence = PySequence_Fast(sequence, "converting sequence.");
+    if (sequence == NULL) {
+        goto finish;
+    }
+
+    Py_ssize_t length = PySequence_Fast_GET_SIZE(sequence);
+    for (Py_ssize_t i = 0; i < length; i++) {
+        PyObject *key_val = PySequence_Fast_GET_ITEM(sequence, i);
+        if (!PyTuple_CheckExact(key_val) || PyTuple_GET_SIZE(key_val) != 2) {
+            PyErr_SetString(PyExc_TypeError, "bad key-value pair.");
+            goto finish;
+        }
+        PyObject *key = PyTuple_GET_ITEM(key_val, 0);
+        PyObject *value = PyTuple_GET_ITEM(key_val, 1);
+        if (!PyTuple_CheckExact(key) || PyTuple_GET_SIZE(key) != key_len) {
+            PyErr_SetString(PyExc_TypeError, "bad key tuple.");
+            goto finish;
+        }
+
+        PyObject *keys[NPY_MAXARGS];
+        for (int j = 0; j < key_len; j++) {
+            keys[j] = PyTuple_GET_ITEM(key, j);
+        }
+        if (i != length - 1) {
+            if (PyArrayIdentityHash_SetItem(tb, keys, value, replace) < 0) {
+                goto finish;
+            }
+        }
+        else {
+            result = PyArrayIdentityHash_GetItem(tb, keys);
+            if (result == NULL) {
+                result = Py_None;
+            }
+            Py_INCREF(result);
+        }
+    }
+
+  finish:
+    Py_DECREF(sequence);
+    PyArrayIdentityHash_Dealloc(tb);
+    return result;
+}
+
+
+/*
  * Test C-api level item getting.
  */
 static PyObject *
@@ -2345,6 +2434,9 @@ static PyMethodDef Multiarray_TestsMethods[] = {
         "Return a list with info on all available casts. Some of the info"
         "may differ for an actual cast if it uses value-based casting "
         "(flexible types)."},
+    {"identityhash_tester",
+        (PyCFunction)identityhash_tester,
+        METH_KEYWORDS | METH_FASTCALL, NULL},
     {"array_indexing",
         array_indexing,
         METH_VARARGS, NULL},
diff --git a/numpy/core/tests/test_hashtable.py b/numpy/core/tests/test_hashtable.py
new file mode 100644
index 000000000..bace4c051
--- /dev/null
+++ b/numpy/core/tests/test_hashtable.py
@@ -0,0 +1,30 @@
+import pytest
+
+import random
+from numpy.core._multiarray_tests import identityhash_tester
+
+
+@pytest.mark.parametrize("key_length", [1, 3, 6])
+@pytest.mark.parametrize("length", [1, 16, 2000])
+def test_identity_hashtable(key_length, length):
+    # use a 30 object pool for everything (duplicates will happen)
+    pool = [object() for i in range(20)]
+    keys_vals = []
+    for i in range(length):
+        keys = tuple(random.choices(pool, k=key_length))
+        keys_vals.append((keys, random.choice(pool)))
+
+    dictionary = dict(keys_vals)
+
+    # add a random item at the end:
+    keys_vals.append(random.choice(keys_vals))
+    # the expected one could be different with duplicates:
+    expected = dictionary[keys_vals[-1][0]]
+
+    res = identityhash_tester(key_length, keys_vals, replace=True)
+    assert res is expected
+
+    # check that ensuring one duplicate definitely raises:
+    keys_vals.insert(0, keys_vals[-2])
+    with pytest.raises(RuntimeError):
+        identityhash_tester(key_length, keys_vals)