Repository: numpy/numpy
Pull request: BUG: ``PyDataMem_SetHandler`` check capsule name
<!--         ----------------------------------------------------------------
                MAKE SURE YOUR PR GETS THE ATTENTION IT DESERVES!
                ----------------------------------------------------------------

*  FORMAT IT RIGHT:
      https://www.numpy.org/devdocs/dev/development_workflow.html#writing-the-commit-message

*  IF IT'S A NEW FEATURE OR API CHANGE, TEST THE WATERS:
      https://www.numpy.org/devdocs/dev/development_workflow.html#get-the-mailing-list-s-opinion

*  HIT ALL THE GUIDELINES:
      https://numpy.org/devdocs/dev/index.html#guidelines

*  WHAT TO DO IF WE HAVEN'T GOTTEN BACK TO YOU:
      https://www.numpy.org/devdocs/dev/development_workflow.html#getting-your-pr-reviewed
-->

Closes https://github.com/numpy/numpy/issues/26137.

Tests capsule name and sets `PyErr` if not valid:
```c
    if (!PyCapsule_IsValid(handler, MEM_HANDLER_CAPSULE_NAME)) {
        PyErr_SetString(PyExc_ValueError, "Capsule must be named 'mem_handler'")
        return NULL;
    }
```

Problem statement:
ENH: `PyDataMem_SetHandler` should check whether the capsule is valid
The capsule must be named `mem_handler` (or nothing will work), but that is not checked at registration time.

It would be good to check at registration time to help users who accidentally misname it (I.e. me ;)).

Hints from issue discussion:

Commit message: TST: PyDataMem_SetHandler mem_handler check lint

Patch to review:
diff --git a/numpy/_core/src/multiarray/alloc.c b/numpy/_core/src/multiarray/alloc.c
index 2f3c82bc5909..df64a13a26e8 100644
--- a/numpy/_core/src/multiarray/alloc.c
+++ b/numpy/_core/src/multiarray/alloc.c
@@ -354,7 +354,8 @@ NPY_NO_EXPORT void *
 PyDataMem_UserNEW(size_t size, PyObject *mem_handler)
 {
     void *result;
-    PyDataMem_Handler *handler = (PyDataMem_Handler *) PyCapsule_GetPointer(mem_handler, "mem_handler");
+    PyDataMem_Handler *handler = (PyDataMem_Handler *) PyCapsule_GetPointer(
+            mem_handler, MEM_HANDLER_CAPSULE_NAME);
     if (handler == NULL) {
         return NULL;
     }
@@ -368,7 +369,8 @@ NPY_NO_EXPORT void *
 PyDataMem_UserNEW_ZEROED(size_t nmemb, size_t size, PyObject *mem_handler)
 {
     void *result;
-    PyDataMem_Handler *handler = (PyDataMem_Handler *) PyCapsule_GetPointer(mem_handler, "mem_handler");
+    PyDataMem_Handler *handler = (PyDataMem_Handler *) PyCapsule_GetPointer(
+            mem_handler, MEM_HANDLER_CAPSULE_NAME);
     if (handler == NULL) {
         return NULL;
     }
@@ -381,7 +383,8 @@ PyDataMem_UserNEW_ZEROED(size_t nmemb, size_t size, PyObject *mem_handler)
 NPY_NO_EXPORT void
 PyDataMem_UserFREE(void *ptr, size_t size, PyObject *mem_handler)
 {
-    PyDataMem_Handler *handler = (PyDataMem_Handler *) PyCapsule_GetPointer(mem_handler, "mem_handler");
+    PyDataMem_Handler *handler = (PyDataMem_Handler *) PyCapsule_GetPointer(
+            mem_handler, MEM_HANDLER_CAPSULE_NAME);
     if (handler == NULL) {
         WARN_NO_RETURN(PyExc_RuntimeWarning,
                      "Could not get pointer to 'mem_handler' from PyCapsule");
@@ -395,7 +398,8 @@ NPY_NO_EXPORT void *
 PyDataMem_UserRENEW(void *ptr, size_t size, PyObject *mem_handler)
 {
     void *result;
-    PyDataMem_Handler *handler = (PyDataMem_Handler *) PyCapsule_GetPointer(mem_handler, "mem_handler");
+    PyDataMem_Handler *handler = (PyDataMem_Handler *) PyCapsule_GetPointer(
+            mem_handler, MEM_HANDLER_CAPSULE_NAME);
     if (handler == NULL) {
         return NULL;
     }
@@ -427,6 +431,10 @@ PyDataMem_SetHandler(PyObject *handler)
     if (handler == NULL) {
         handler = PyDataMem_DefaultHandler;
     }
+    if (!PyCapsule_IsValid(handler, MEM_HANDLER_CAPSULE_NAME)) {
+        PyErr_SetString(PyExc_ValueError, "Capsule must be named 'mem_handler'");
+        return NULL;
+    }
     token = PyContextVar_Set(current_handler, handler);
     if (token == NULL) {
         Py_DECREF(old_handler);
@@ -477,7 +485,8 @@ get_handler_name(PyObject *NPY_UNUSED(self), PyObject *args)
             return NULL;
         }
     }
-    handler = (PyDataMem_Handler *) PyCapsule_GetPointer(mem_handler, "mem_handler");
+    handler = (PyDataMem_Handler *) PyCapsule_GetPointer(
+            mem_handler, MEM_HANDLER_CAPSULE_NAME);
     if (handler == NULL) {
         Py_DECREF(mem_handler);
         return NULL;
@@ -514,7 +523,8 @@ get_handler_version(PyObject *NPY_UNUSED(self), PyObject *args)
             return NULL;
         }
     }
-    handler = (PyDataMem_Handler *) PyCapsule_GetPointer(mem_handler, "mem_handler");
+    handler = (PyDataMem_Handler *) PyCapsule_GetPointer(
+            mem_handler, MEM_HANDLER_CAPSULE_NAME);
     if (handler == NULL) {
         Py_DECREF(mem_handler);
         return NULL;
diff --git a/numpy/_core/src/multiarray/alloc.h b/numpy/_core/src/multiarray/alloc.h
index 186eb54870ab..aed2095fe73c 100644
--- a/numpy/_core/src/multiarray/alloc.h
+++ b/numpy/_core/src/multiarray/alloc.h
@@ -6,6 +6,7 @@
 #include "numpy/ndarraytypes.h"
 
 #define NPY_TRACE_DOMAIN 389047
+#define MEM_HANDLER_CAPSULE_NAME "mem_handler"
 
 NPY_NO_EXPORT PyObject *
 _get_madvise_hugepage(PyObject *NPY_UNUSED(self), PyObject *NPY_UNUSED(args));
diff --git a/numpy/_core/src/multiarray/multiarraymodule.c b/numpy/_core/src/multiarray/multiarraymodule.c
index 4946465617bc..1cda149136c4 100644
--- a/numpy/_core/src/multiarray/multiarraymodule.c
+++ b/numpy/_core/src/multiarray/multiarraymodule.c
@@ -5234,7 +5234,8 @@ PyMODINIT_FUNC PyInit__multiarray_umath(void) {
     /*
      * Initialize the default PyDataMem_Handler capsule singleton.
      */
-    PyDataMem_DefaultHandler = PyCapsule_New(&default_handler, "mem_handler", NULL);
+    PyDataMem_DefaultHandler = PyCapsule_New(
+            &default_handler, MEM_HANDLER_CAPSULE_NAME, NULL);
     if (PyDataMem_DefaultHandler == NULL) {
         goto err;
     }
diff --git a/numpy/_core/tests/__init__.py b/numpy/_core/tests/__init__.py
deleted file mode 100644
index e69de29bb2d1..000000000000
diff --git a/numpy/_core/tests/test_mem_policy.py b/numpy/_core/tests/test_mem_policy.py
index 3c7f6bb34661..e8559ac0cced 100644
--- a/numpy/_core/tests/test_mem_policy.py
+++ b/numpy/_core/tests/test_mem_policy.py
@@ -46,6 +46,16 @@ def get_module(tmp_path):
              Py_DECREF(secret_data);
              return old;
          """),
+        ("set_wrong_capsule_name_data_policy", "METH_NOARGS", """
+             PyObject *secret_data =
+                 PyCapsule_New(&secret_data_handler, "not_mem_handler", NULL);
+             if (secret_data == NULL) {
+                 return NULL;
+             }
+             PyObject *old = PyDataMem_SetHandler(secret_data);
+             Py_DECREF(secret_data);
+             return old;
+         """),
         ("set_old_policy", "METH_O", """
              PyObject *old;
              if (args != NULL && PyCapsule_CheckExact(args)) {
@@ -252,6 +262,10 @@ def test_set_policy(get_module):
         get_module.set_old_policy(orig_policy)
         assert get_handler_name() == orig_policy_name
 
+    with pytest.raises(ValueError,
+                       match="Capsule must be named 'mem_handler'"):
+        get_module.set_wrong_capsule_name_data_policy()
+
 
 @pytest.mark.skipif(sys.version_info >= (3, 12), reason="no numpy.distutils")
 def test_default_policy_singleton(get_module):


Code review comments:
diff_hunk: @@ -46,6 +46,16 @@ def get_module(tmp_path):
              Py_DECREF(secret_data);
              return old;
          """),
+        ("set_wrong_capsule_name_data_policy", "METH_NOARGS", """
+             PyObject *secret_data =
path: numpy/_core/tests/test_mem_policy.py
text: @user1:
```suggestion
             PyObject *wrong_name_capsule =
```

diff_hunk: @@ -46,6 +46,16 @@ def get_module(tmp_path):
              Py_DECREF(secret_data);
              return old;
          """),
+        ("set_wrong_capsule_name_data_policy", "METH_NOARGS", """
+             PyObject *secret_data =
path: numpy/_core/tests/test_mem_policy.py
text: @author:
I used the same `struct` for the secret_data_policy. Is this ok or should I create a struct just for this test?

@user1:
That's can be anything here really.  But let's rename the object (will just apply the suggestion though).

diff_hunk: @@ -46,6 +46,16 @@ def get_module(tmp_path):
              Py_DECREF(secret_data);
              return old;
          """),
+        ("set_wrong_capsule_name_data_policy", "METH_NOARGS", """
+             PyObject *secret_data =
+                 PyCapsule_New(&secret_data_handler, "not_mem_handler", NULL);
+             if (secret_data == NULL) {
+                 return NULL;
+             }
+             PyObject *old = PyDataMem_SetHandler(secret_data);
+             Py_DECREF(secret_data);
path: numpy/_core/tests/test_mem_policy.py
text: @user1:
```suggestion
             if (wrong_name_capsule == NULL) {
                 return NULL;
             }
             PyObject *old = PyDataMem_SetHandler(wrong_name_capsule);
             Py_DECREF(wrong_name_capsule);
```