Let the argument clinic do the type checking for heapq (GH-20284)

2 files changed, 117 insertions, 46 deletions

diff --git a/Modules/_heapqmodule.c b/Modules/_heapqmodule.c
index 4e85e04..193478d 100644
--- a/Modules/_heapqmodule.c
+++ b/Modules/_heapqmodule.c
@@ -113,7 +113,7 @@ siftup(PyListObject *heap, Py_ssize_t pos)
 /*[clinic input]
 _heapq.heappush
 
-    heap: object
+    heap: object(subclass_of='&PyList_Type')
     item: object
     /
 
@@ -122,13 +122,8 @@ Push item onto heap, maintaining the heap invariant.
 
 static PyObject *
 _heapq_heappush_impl(PyObject *module, PyObject *heap, PyObject *item)
-/*[clinic end generated code: output=912c094f47663935 input=7913545cb5118842]*/
+/*[clinic end generated code: output=912c094f47663935 input=7c69611f3698aceb]*/
 {
-    if (!PyList_Check(heap)) {
-        PyErr_SetString(PyExc_TypeError, "heap argument must be a list");
-        return NULL;
-    }
-
     if (PyList_Append(heap, item))
         return NULL;
 
@@ -143,11 +138,6 @@ heappop_internal(PyObject *heap, int siftup_func(PyListObject *, Py_ssize_t))
     PyObject *lastelt, *returnitem;
     Py_ssize_t n;
 
-    if (!PyList_Check(heap)) {
-        PyErr_SetString(PyExc_TypeError, "heap argument must be a list");
-        return NULL;
-    }
-
     /* raises IndexError if the heap is empty */
     n = PyList_GET_SIZE(heap);
     if (n == 0) {
@@ -177,15 +167,15 @@ heappop_internal(PyObject *heap, int siftup_func(PyListObject *, Py_ssize_t))
 /*[clinic input]
 _heapq.heappop
 
-    heap: object
+    heap: object(subclass_of='&PyList_Type')
     /
 
 Pop the smallest item off the heap, maintaining the heap invariant.
 [clinic start generated code]*/
 
 static PyObject *
-_heapq_heappop(PyObject *module, PyObject *heap)
-/*[clinic end generated code: output=e1bbbc9866bce179 input=9bd36317b806033d]*/
+_heapq_heappop_impl(PyObject *module, PyObject *heap)
+/*[clinic end generated code: output=96dfe82d37d9af76 input=91487987a583c856]*/
 {
     return heappop_internal(heap, siftup);
 }
@@ -195,11 +185,6 @@ heapreplace_internal(PyObject *heap, PyObject *item, int siftup_func(PyListObjec
 {
     PyObject *returnitem;
 
-    if (!PyList_Check(heap)) {
-        PyErr_SetString(PyExc_TypeError, "heap argument must be a list");
-        return NULL;
-    }
-
     if (PyList_GET_SIZE(heap) == 0) {
         PyErr_SetString(PyExc_IndexError, "index out of range");
         return NULL;
@@ -219,7 +204,7 @@ heapreplace_internal(PyObject *heap, PyObject *item, int siftup_func(PyListObjec
 /*[clinic input]
 _heapq.heapreplace
 
-    heap: object
+    heap: object(subclass_of='&PyList_Type')
     item: object
     /
 
@@ -236,7 +221,7 @@ this routine unless written as part of a conditional replacement:
 
 static PyObject *
 _heapq_heapreplace_impl(PyObject *module, PyObject *heap, PyObject *item)
-/*[clinic end generated code: output=82ea55be8fbe24b4 input=e57ae8f4ecfc88e3]*/
+/*[clinic end generated code: output=82ea55be8fbe24b4 input=719202ac02ba10c8]*/
 {
     return heapreplace_internal(heap, item, siftup);
 }
@@ -244,7 +229,7 @@ _heapq_heapreplace_impl(PyObject *module, PyObject *heap, PyObject *item)
 /*[clinic input]
 _heapq.heappushpop
 
-    heap: object
+    heap: object(subclass_of='&PyList_Type')
     item: object
     /
 
@@ -256,16 +241,11 @@ a separate call to heappop().
 
 static PyObject *
 _heapq_heappushpop_impl(PyObject *module, PyObject *heap, PyObject *item)
-/*[clinic end generated code: output=67231dc98ed5774f input=eb48c90ba77b2214]*/
+/*[clinic end generated code: output=67231dc98ed5774f input=5dc701f1eb4a4aa7]*/
 {
     PyObject *returnitem;
     int cmp;
 
-    if (!PyList_Check(heap)) {
-        PyErr_SetString(PyExc_TypeError, "heap argument must be a list");
-        return NULL;
-    }
-
     if (PyList_GET_SIZE(heap) == 0) {
         Py_INCREF(item);
         return item;
@@ -367,11 +347,6 @@ heapify_internal(PyObject *heap, int siftup_func(PyListObject *, Py_ssize_t))
 {
     Py_ssize_t i, n;
 
-    if (!PyList_Check(heap)) {
-        PyErr_SetString(PyExc_TypeError, "heap argument must be a list");
-        return NULL;
-    }
-
     /* For heaps likely to be bigger than L1 cache, we use the cache
        friendly heapify function.  For smaller heaps that fit entirely
        in cache, we prefer the simpler algorithm with less branching.
@@ -396,15 +371,15 @@ heapify_internal(PyObject *heap, int siftup_func(PyListObject *, Py_ssize_t))
 /*[clinic input]
 _heapq.heapify
 
-    heap: object
+    heap: object(subclass_of='&PyList_Type')
     /
 
 Transform list into a heap, in-place, in O(len(heap)) time.
 [clinic start generated code]*/
 
 static PyObject *
-_heapq_heapify(PyObject *module, PyObject *heap)
-/*[clinic end generated code: output=11483f23627c4616 input=872c87504b8de970]*/
+_heapq_heapify_impl(PyObject *module, PyObject *heap)
+/*[clinic end generated code: output=e63a636fcf83d6d0 input=53bb7a2166febb73]*/
 {
     return heapify_internal(heap, siftup);
 }
@@ -508,15 +483,15 @@ siftup_max(PyListObject *heap, Py_ssize_t pos)
 /*[clinic input]
 _heapq._heappop_max
 
-    heap: object
+    heap: object(subclass_of='&PyList_Type')
     /
 
 Maxheap variant of heappop.
 [clinic start generated code]*/
 
 static PyObject *
-_heapq__heappop_max(PyObject *module, PyObject *heap)
-/*[clinic end generated code: output=acd30acf6384b13c input=62ede3ba9117f541]*/
+_heapq__heappop_max_impl(PyObject *module, PyObject *heap)
+/*[clinic end generated code: output=9e77aadd4e6a8760 input=362c06e1c7484793]*/
 {
     return heappop_internal(heap, siftup_max);
 }
@@ -524,7 +499,7 @@ _heapq__heappop_max(PyObject *module, PyObject *heap)
 /*[clinic input]
 _heapq._heapreplace_max
 
-    heap: object
+    heap: object(subclass_of='&PyList_Type')
     item: object
     /
 
@@ -534,7 +509,7 @@ Maxheap variant of heapreplace.
 static PyObject *
 _heapq__heapreplace_max_impl(PyObject *module, PyObject *heap,
                              PyObject *item)
-/*[clinic end generated code: output=8ad7545e4a5e8adb input=6d8f25131e0f0e5f]*/
+/*[clinic end generated code: output=8ad7545e4a5e8adb input=f2dd27cbadb948d7]*/
 {
     return heapreplace_internal(heap, item, siftup_max);
 }
@@ -542,15 +517,15 @@ _heapq__heapreplace_max_impl(PyObject *module, PyObject *heap,
 /*[clinic input]
 _heapq._heapify_max
 
-    heap: object
+    heap: object(subclass_of='&PyList_Type')
     /
 
 Maxheap variant of heapify.
 [clinic start generated code]*/
 
 static PyObject *
-_heapq__heapify_max(PyObject *module, PyObject *heap)
-/*[clinic end generated code: output=1c6bb6b60d6a2133 input=cdfcc6835b14110d]*/
+_heapq__heapify_max_impl(PyObject *module, PyObject *heap)
+/*[clinic end generated code: output=2cb028beb4a8b65e input=c1f765ee69f124b8]*/
 {
     return heapify_internal(heap, siftup_max);
 }
diff --git a/Modules/clinic/_heapqmodule.c.h b/Modules/clinic/_heapqmodule.c.h
index 5540370..a894315 100644
--- a/Modules/clinic/_heapqmodule.c.h
+++ b/Modules/clinic/_heapqmodule.c.h
@@ -24,6 +24,10 @@ _heapq_heappush(PyObject *module, PyObject *const *args, Py_ssize_t nargs)
     if (!_PyArg_CheckPositional("heappush", nargs, 2, 2)) {
         goto exit;
     }
+    if (!PyList_Check(args[0])) {
+        _PyArg_BadArgument("heappush", "argument 1", "list", args[0]);
+        goto exit;
+    }
     heap = args[0];
     item = args[1];
     return_value = _heapq_heappush_impl(module, heap, item);
@@ -41,6 +45,26 @@ PyDoc_STRVAR(_heapq_heappop__doc__,
 #define _HEAPQ_HEAPPOP_METHODDEF    \
     {"heappop", (PyCFunction)_heapq_heappop, METH_O, _heapq_heappop__doc__},
 
+static PyObject *
+_heapq_heappop_impl(PyObject *module, PyObject *heap);
+
+static PyObject *
+_heapq_heappop(PyObject *module, PyObject *arg)
+{
+    PyObject *return_value = NULL;
+    PyObject *heap;
+
+    if (!PyList_Check(arg)) {
+        _PyArg_BadArgument("heappop", "argument", "list", arg);
+        goto exit;
+    }
+    heap = arg;
+    return_value = _heapq_heappop_impl(module, heap);
+
+exit:
+    return return_value;
+}
+
 PyDoc_STRVAR(_heapq_heapreplace__doc__,
 "heapreplace($module, heap, item, /)\n"
 "--\n"
@@ -71,6 +95,10 @@ _heapq_heapreplace(PyObject *module, PyObject *const *args, Py_ssize_t nargs)
     if (!_PyArg_CheckPositional("heapreplace", nargs, 2, 2)) {
         goto exit;
     }
+    if (!PyList_Check(args[0])) {
+        _PyArg_BadArgument("heapreplace", "argument 1", "list", args[0]);
+        goto exit;
+    }
     heap = args[0];
     item = args[1];
     return_value = _heapq_heapreplace_impl(module, heap, item);
@@ -104,6 +132,10 @@ _heapq_heappushpop(PyObject *module, PyObject *const *args, Py_ssize_t nargs)
     if (!_PyArg_CheckPositional("heappushpop", nargs, 2, 2)) {
         goto exit;
     }
+    if (!PyList_Check(args[0])) {
+        _PyArg_BadArgument("heappushpop", "argument 1", "list", args[0]);
+        goto exit;
+    }
     heap = args[0];
     item = args[1];
     return_value = _heapq_heappushpop_impl(module, heap, item);
@@ -121,6 +153,26 @@ PyDoc_STRVAR(_heapq_heapify__doc__,
 #define _HEAPQ_HEAPIFY_METHODDEF    \
     {"heapify", (PyCFunction)_heapq_heapify, METH_O, _heapq_heapify__doc__},
 
+static PyObject *
+_heapq_heapify_impl(PyObject *module, PyObject *heap);
+
+static PyObject *
+_heapq_heapify(PyObject *module, PyObject *arg)
+{
+    PyObject *return_value = NULL;
+    PyObject *heap;
+
+    if (!PyList_Check(arg)) {
+        _PyArg_BadArgument("heapify", "argument", "list", arg);
+        goto exit;
+    }
+    heap = arg;
+    return_value = _heapq_heapify_impl(module, heap);
+
+exit:
+    return return_value;
+}
+
 PyDoc_STRVAR(_heapq__heappop_max__doc__,
 "_heappop_max($module, heap, /)\n"
 "--\n"
@@ -130,6 +182,26 @@ PyDoc_STRVAR(_heapq__heappop_max__doc__,
 #define _HEAPQ__HEAPPOP_MAX_METHODDEF    \
     {"_heappop_max", (PyCFunction)_heapq__heappop_max, METH_O, _heapq__heappop_max__doc__},
 
+static PyObject *
+_heapq__heappop_max_impl(PyObject *module, PyObject *heap);
+
+static PyObject *
+_heapq__heappop_max(PyObject *module, PyObject *arg)
+{
+    PyObject *return_value = NULL;
+    PyObject *heap;
+
+    if (!PyList_Check(arg)) {
+        _PyArg_BadArgument("_heappop_max", "argument", "list", arg);
+        goto exit;
+    }
+    heap = arg;
+    return_value = _heapq__heappop_max_impl(module, heap);
+
+exit:
+    return return_value;
+}
+
 PyDoc_STRVAR(_heapq__heapreplace_max__doc__,
 "_heapreplace_max($module, heap, item, /)\n"
 "--\n"
@@ -153,6 +225,10 @@ _heapq__heapreplace_max(PyObject *module, PyObject *const *args, Py_ssize_t narg
     if (!_PyArg_CheckPositional("_heapreplace_max", nargs, 2, 2)) {
         goto exit;
     }
+    if (!PyList_Check(args[0])) {
+        _PyArg_BadArgument("_heapreplace_max", "argument 1", "list", args[0]);
+        goto exit;
+    }
     heap = args[0];
     item = args[1];
     return_value = _heapq__heapreplace_max_impl(module, heap, item);
@@ -169,4 +245,24 @@ PyDoc_STRVAR(_heapq__heapify_max__doc__,
 
 #define _HEAPQ__HEAPIFY_MAX_METHODDEF    \
     {"_heapify_max", (PyCFunction)_heapq__heapify_max, METH_O, _heapq__heapify_max__doc__},
-/*[clinic end generated code: output=37ef2a3319971c8d input=a9049054013a1b77]*/
+
+static PyObject *
+_heapq__heapify_max_impl(PyObject *module, PyObject *heap);
+
+static PyObject *
+_heapq__heapify_max(PyObject *module, PyObject *arg)
+{
+    PyObject *return_value = NULL;
+    PyObject *heap;
+
+    if (!PyList_Check(arg)) {
+        _PyArg_BadArgument("_heapify_max", "argument", "list", arg);
+        goto exit;
+    }
+    heap = arg;
+    return_value = _heapq__heapify_max_impl(module, heap);
+
+exit:
+    return return_value;
+}
+/*[clinic end generated code: output=9975cf51762878d5 input=a9049054013a1b77]*/