Fix and improve Objects/lazyimportobject.c (#31)

* re-order imports

* simplify static type definition

* fix undefined behaviors

* reject keyword arguments for `lazy_import`

Author: picnixz

Objects/lazyimportobject.c

@@ -1,11 +1,14 @@
 /* Lazy object implementation */
 
 #include "Python.h"
-#include "pycore_import.h"
-#include "pycore_lazyimportobject.h"
-#include "pycore_frame.h"
 #include "pycore_ceval.h"
+#include "pycore_frame.h"
+#include "pycore_import.h"
 #include "pycore_interpframe.h"
+#include "pycore_lazyimportobject.h"
+#include "pycore_modsupport.h"
+
+#define PyLazyImportObject_CAST(op) ((PyLazyImportObject *)(op))
 
 PyObject *
 _PyLazyImport_New(PyObject *builtins, PyObject *from, PyObject *attr)
@@ -46,8 +49,9 @@ _PyLazyImport_New(PyObject *builtins, PyObject *from, PyObject *attr)
 }
 
 static int
-lazy_import_traverse(PyLazyImportObject *m, visitproc visit, void *arg)
+lazy_import_traverse(PyObject *op, visitproc visit, void *arg)
 {
+    PyLazyImportObject *m = PyLazyImportObject_CAST(op);
     Py_VISIT(m->lz_builtins);
     Py_VISIT(m->lz_from);
     Py_VISIT(m->lz_attr);
@@ -56,8 +60,9 @@ lazy_import_traverse(PyLazyImportObject *m, visitproc visit, void *arg)
 }
 
 static int
-lazy_import_clear(PyLazyImportObject *m)
+lazy_import_clear(PyObject *op)
 {
+    PyLazyImportObject *m = PyLazyImportObject_CAST(op);
     Py_CLEAR(m->lz_builtins);
     Py_CLEAR(m->lz_from);
     Py_CLEAR(m->lz_attr);
@@ -66,11 +71,11 @@ lazy_import_clear(PyLazyImportObject *m)
 }
 
 static void
-lazy_import_dealloc(PyLazyImportObject *m)
+lazy_import_dealloc(PyObject *op)
 {
-    _PyObject_GC_UNTRACK(m);
-    lazy_import_clear(m);
-    Py_TYPE(m)->tp_free((PyObject *)m);
+    _PyObject_GC_UNTRACK(op);
+    (void)lazy_import_clear(op);
+    Py_TYPE(op)->tp_free(op);
 }
 
 static PyObject *
@@ -79,60 +84,69 @@ lazy_import_name(PyLazyImportObject *m)
     if (m->lz_attr != NULL) {
         if (PyUnicode_Check(m->lz_attr)) {
             return PyUnicode_FromFormat("%U.%U", m->lz_from, m->lz_attr);
-        } else {
+        }
+        else {
             return PyUnicode_FromFormat("%U...", m->lz_from);
         }
     }
-    Py_INCREF(m->lz_from);
-    return m->lz_from;
+    return Py_NewRef(m->lz_from);
 }
 
 static PyObject *
-lazy_import_repr(PyLazyImportObject *m)
+lazy_import_repr(PyObject *op)
 {
+    PyLazyImportObject *m = PyLazyImportObject_CAST(op);
     PyObject *name = lazy_import_name(m);
     if (name == NULL) {
         return NULL;
     }
-    PyObject *res = PyUnicode_FromFormat("<lazy_import '%U'>", name);
+    PyObject *res = PyUnicode_FromFormat("<%T '%U'>", op, name);
     Py_DECREF(name);
     return res;
 }
 
 static PyObject *
 lazy_import_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
 {
-    if (PyTuple_GET_SIZE(args) != 2 && PyTuple_GET_SIZE(args) != 3) {
-        PyErr_SetString(PyExc_ValueError, "lazy_import expected 2-3 arguments");
+    PyTypeObject *base_tp = &PyLazyImport_Type;
+    if (
+        (type == base_tp || type->tp_init == base_tp->tp_init)
+        && !_PyArg_NoKeywords("lazy_import", kwds)
+    ) {
         return NULL;
     }
 
-    PyObject *builtins = PyTuple_GET_ITEM(args, 0);
-    PyObject *from = PyTuple_GET_ITEM(args, 1);
-    PyObject *attr = NULL;
-    if (PyTuple_GET_SIZE(args) == 3) {
-        attr = PyTuple_GET_ITEM(args, 2);
+    Py_ssize_t nargs = PyTuple_GET_SIZE(args);
+    if (!_PyArg_CheckPositional("lazy_import", nargs, 2, 3)) {
+        return NULL;
     }
 
+    PyObject *builtins = PyTuple_GET_ITEM(args, 0);
+    PyObject *from = PyTuple_GET_ITEM(args, 1);
+    PyObject *attr = nargs == 3 ? PyTuple_GET_ITEM(args, 2) : NULL;
     return _PyLazyImport_New(builtins, from, attr);
 }
 
 PyObject *
-_PyLazyImport_GetName(PyObject *lazy_import)
+_PyLazyImport_GetName(PyObject *op)
 {
+    PyLazyImportObject *lazy_import = PyLazyImportObject_CAST(op);
     assert(PyLazyImport_CheckExact(lazy_import));
-    return lazy_import_name((PyLazyImportObject *)lazy_import);
+    return lazy_import_name(lazy_import);
 }
 
 static PyObject *
-lazy_resolve(PyObject *self, PyObject *args)
+lazy_import_resolve(PyObject *self, PyObject *args)
 {
     return _PyImport_LoadLazyImportTstate(PyThreadState_GET(), self);
 }
 
-static PyMethodDef lazy_methods[] = {
-    {"resolve", lazy_resolve, METH_NOARGS, PyDoc_STR("resolves the lazy import and returns the actual object")},
-    {0}
+static PyMethodDef lazy_import_methods[] = {
+    {
+        "resolve", lazy_import_resolve, METH_NOARGS,
+        PyDoc_STR("resolves the lazy import and returns the actual object")
+    },
+    {NULL, NULL}
 };
 
 
@@ -147,43 +161,16 @@ PyDoc_STRVAR(lazy_import_doc,
 
 PyTypeObject PyLazyImport_Type = {
     PyVarObject_HEAD_INIT(&PyType_Type, 0)
-    "lazy_import",                              /* tp_name */
-    sizeof(PyLazyImportObject),                 /* tp_basicsize */
-    0,                                          /* tp_itemsize */
-    (destructor)lazy_import_dealloc,            /* tp_dealloc */
-    0,                                          /* tp_print */
-    0,                                          /* tp_getattr */
-    0,                                          /* tp_setattr */
-    0,                                          /* tp_reserved */
-    (reprfunc)lazy_import_repr,                 /* tp_repr */
-    0,                                          /* tp_as_number */
-    0,                                          /* tp_as_sequence */
-    0,                                          /* tp_as_mapping */
-    0,                                          /* tp_hash */
-    0,                                          /* tp_call */
-    0,                                          /* tp_str */
-    0,                                          /* tp_getattro */
-    0,                                          /* tp_setattro */
-    0,                                          /* tp_as_buffer */
-    Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
-        Py_TPFLAGS_BASETYPE,                    /* tp_flags */
-    lazy_import_doc,                            /* tp_doc */
-    (traverseproc)lazy_import_traverse,         /* tp_traverse */
-    (inquiry)lazy_import_clear,                 /* tp_clear */
-    0,                                          /* tp_richcompare */
-    0,                                          /* tp_weaklistoffset */
-    0,                                          /* tp_iter */
-    0,                                          /* tp_iternext */
-    lazy_methods,                               /* tp_methods */
-    0,                                          /* tp_members */
-    0,                                          /* tp_getset */
-    0,                                          /* tp_base */
-    0,                                          /* tp_dict */
-    0,                                          /* tp_descr_get */
-    0,                                          /* tp_descr_set */
-    0,                                          /* tp_dictoffset */
-    0,                                          /* tp_init */
-    PyType_GenericAlloc,                        /* tp_alloc */
-    lazy_import_new,                            /* tp_new */
-    PyObject_GC_Del,                            /* tp_free */
+    .tp_name = "lazy_import",
+    .tp_basicsize = sizeof(PyLazyImportObject),
+    .tp_dealloc = lazy_import_dealloc,
+    .tp_repr = lazy_import_repr,
+    .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_BASETYPE,
+    .tp_doc = lazy_import_doc,
+    .tp_traverse = lazy_import_traverse,
+    .tp_clear = lazy_import_clear,
+    .tp_methods = lazy_import_methods,
+    .tp_alloc = PyType_GenericAlloc,
+    .tp_new = lazy_import_new,
+    .tp_free = PyObject_GC_Del,
 };