gh-87390: Fix starred tuple equality and pickling (GH-92337)

python · May 5, 2022 · 1ed8d03 · 1ed8d03
1 parent 354ab7a
commit 1ed8d03
Show file tree

Hide file tree

Showing 4 changed files with 31 additions and 0 deletions.
diff --git a/Include/internal/pycore_global_strings.h b/Include/internal/pycore_global_strings.h
@@ -300,6 +300,7 @@ struct _Py_global_strings {
         STRUCT_FOR_ID(n_unnamed_fields)
         STRUCT_FOR_ID(name)
         STRUCT_FOR_ID(newlines)
+        STRUCT_FOR_ID(next)
         STRUCT_FOR_ID(obj)
         STRUCT_FOR_ID(offset)
         STRUCT_FOR_ID(onceregistry)

diff --git a/Include/internal/pycore_runtime_init.h b/Include/internal/pycore_runtime_init.h
@@ -923,6 +923,7 @@ extern "C" {
                 INIT_ID(n_unnamed_fields), \
                 INIT_ID(name), \
                 INIT_ID(newlines), \
+                INIT_ID(next), \
                 INIT_ID(obj), \
                 INIT_ID(offset), \
                 INIT_ID(onceregistry), \

diff --git a/Lib/test/test_genericalias.py b/Lib/test/test_genericalias.py
@@ -358,6 +358,8 @@ def test_equality(self):
         self.assertNotEqual(dict[str, int], dict[str, str])
         self.assertNotEqual(list, list[int])
         self.assertNotEqual(list[int], list)
+        self.assertNotEqual(list[int], tuple[int])
+        self.assertNotEqual((*tuple[int],)[0], tuple[int])
 
     def test_isinstance(self):
         self.assertTrue(isinstance([], list))
@@ -394,6 +396,7 @@ def test_pickle(self):
                     self.assertEqual(loaded.__origin__, alias.__origin__)
                     self.assertEqual(loaded.__args__, alias.__args__)
                     self.assertEqual(loaded.__parameters__, alias.__parameters__)
+                    self.assertEqual(type(loaded), type(alias))
 
     def test_copy(self):
         class X(list):

diff --git a/Objects/genericaliasobject.c b/Objects/genericaliasobject.c
@@ -567,6 +567,9 @@ ga_richcompare(PyObject *a, PyObject *b, int op)
 
     gaobject *aa = (gaobject *)a;
     gaobject *bb = (gaobject *)b;
+    if (aa->starred != bb->starred) {
+        Py_RETURN_FALSE;
+    }
     int eq = PyObject_RichCompareBool(aa->origin, bb->origin, Py_EQ);
     if (eq < 0) {
         return NULL;
@@ -604,6 +607,16 @@ static PyObject *
 ga_reduce(PyObject *self, PyObject *Py_UNUSED(ignored))
 {
     gaobject *alias = (gaobject *)self;
+    if (alias->starred) {
+        PyObject *tmp = Py_GenericAlias(alias->origin, alias->args);
+        if (tmp != NULL) {
+            Py_SETREF(tmp, PyObject_GetIter(tmp));
+        }
+        if (tmp == NULL) {
+            return NULL;
+        }
+        return Py_BuildValue("N(N)", _PyEval_GetBuiltin(&_Py_ID(next)), tmp);
+    }
     return Py_BuildValue("O(OO)", Py_TYPE(alias),
                          alias->origin, alias->args);
 }
@@ -775,6 +788,18 @@ ga_iter_clear(PyObject *self) {
     return 0;
 }
 
+static PyObject *
+ga_iter_reduce(PyObject *self, PyObject *Py_UNUSED(ignored))
+{
+    gaiterobject *gi = (gaiterobject *)self;
+    return Py_BuildValue("N(O)", _PyEval_GetBuiltin(&_Py_ID(iter)), gi->obj);
+}
+
+static PyMethodDef ga_iter_methods[] = {
+    {"__reduce__", ga_iter_reduce, METH_NOARGS},
+    {0}
+};
+
 // gh-91632: _Py_GenericAliasIterType is exported  to be cleared
 // in _PyTypes_FiniTypes.
 PyTypeObject _Py_GenericAliasIterType = {
@@ -784,6 +809,7 @@ PyTypeObject _Py_GenericAliasIterType = {
     .tp_iter = PyObject_SelfIter,
     .tp_iternext = (iternextfunc)ga_iternext,
     .tp_traverse = (traverseproc)ga_iter_traverse,
+    .tp_methods = ga_iter_methods,
     .tp_dealloc = (destructor)ga_iter_dealloc,
     .tp_clear = (inquiry)ga_iter_clear,
     .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,