bpo-43224: Implement PEP 646 changes to genericaliasobject.c (GH-31019)

mrahtz · JelleZijlstra · kumaraditya303 · web-flow · commit af2277e461ae · 2022-03-12T20:20:12.000+08:00
Specifically, prepare for starring of tuples via a new genericalias iter type. GenericAlias also partially supports the iterator protocol after this change.

Co-authored-by: Jelle Zijlstra &lt;jelle.zijlstra@gmail.com&gt;
Co-authored-by: Kumar Aditya &lt;59607654+kumaraditya303@users.noreply.github.com&gt;
Co-authored-by: Ken Jin &lt;28750310+Fidget-Spinner@users.noreply.github.com&gt;
diff --git a/Lib/test/test_genericalias.py b/Lib/test/test_genericalias.py
@@ -169,6 +169,24 @@ class MyList(list):
         self.assertEqual(repr(list[str]), 'list[str]')
         self.assertEqual(repr(list[()]), 'list[()]')
         self.assertEqual(repr(tuple[int, ...]), 'tuple[int, ...]')
+        x1 = tuple[
+            tuple(  # Effectively the same as starring; TODO
+                tuple[int]
+            )
+        ]
+        self.assertEqual(repr(x1), 'tuple[*tuple[int]]')
+        x2 = tuple[
+            tuple(  # Ditto TODO
+                tuple[int, str]
+            )
+        ]
+        self.assertEqual(repr(x2), 'tuple[*tuple[int, str]]')
+        x3 = tuple[
+            tuple(  # Ditto TODO
+                tuple[int, ...]
+            )
+        ]
+        self.assertEqual(repr(x3), 'tuple[*tuple[int, ...]]')
         self.assertTrue(repr(MyList[int]).endswith('.BaseTest.test_repr.<locals>.MyList[int]'))
         self.assertEqual(repr(list[str]()), '[]')  # instances should keep their normal repr
 
@@ -182,6 +200,7 @@ def test_exposed_type(self):
 
     def test_parameters(self):
         from typing import List, Dict, Callable
+
         D0 = dict[str, int]
         self.assertEqual(D0.__args__, (str, int))
         self.assertEqual(D0.__parameters__, ())
@@ -197,6 +216,7 @@ def test_parameters(self):
         D2b = dict[T, T]
         self.assertEqual(D2b.__args__, (T, T))
         self.assertEqual(D2b.__parameters__, (T,))
+
         L0 = list[str]
         self.assertEqual(L0.__args__, (str,))
         self.assertEqual(L0.__parameters__, ())
@@ -219,6 +239,45 @@ def test_parameters(self):
         self.assertEqual(L5.__args__, (Callable[[K, V], K],))
         self.assertEqual(L5.__parameters__, (K, V))
 
+        T1 = tuple[
+            tuple(  # Ditto TODO
+                tuple[int]
+            )
+        ]
+        self.assertEqual(
+            T1.__args__,
+            tuple(  # Ditto TODO
+                tuple[int]
+            )
+        )
+        self.assertEqual(T1.__parameters__, ())
+
+        T2 = tuple[
+            tuple(  # Ditto TODO
+                tuple[T]
+            )
+        ]
+        self.assertEqual(
+            T2.__args__,
+            tuple(  # Ditto TODO
+                tuple[T]
+            )
+        )
+        self.assertEqual(T2.__parameters__, (T,))
+
+        T4 = tuple[
+            tuple(  # Ditto TODO
+                tuple[int, str]
+            )
+        ]
+        self.assertEqual(
+            T4.__args__,
+            tuple(  # Ditto TODO
+                tuple[int, str]
+            )
+        )
+        self.assertEqual(T4.__parameters__, ())
+
     def test_parameter_chaining(self):
         from typing import List, Dict, Union, Callable
         self.assertEqual(list[T][int], list[int])
@@ -249,6 +308,19 @@ def test_parameter_chaining(self):
     def test_equality(self):
         self.assertEqual(list[int], list[int])
         self.assertEqual(dict[str, int], dict[str, int])
+        self.assertEqual((*tuple[int],)[0], (*tuple[int],)[0])
+        self.assertEqual(
+            tuple[
+                tuple(  # Effectively the same as starring; TODO
+                    tuple[int]
+                )
+            ],
+            tuple[
+                tuple(  # Ditto TODO
+                    tuple[int]
+                )
+            ]
+        )
         self.assertNotEqual(dict[str, int], dict[str, str])
         self.assertNotEqual(list, list[int])
         self.assertNotEqual(list[int], list)
@@ -346,6 +418,24 @@ def __new__(cls, *args, **kwargs):
         with self.assertRaises(TypeError):
             Bad(list, int, bad=int)
 
+    def test_iter_creates_starred_tuple(self):
+        t = tuple[int, str]
+        iter_t = iter(t)
+        x = next(iter_t)
+        self.assertEqual(repr(x), '*tuple[int, str]')
+
+    def test_calling_next_twice_raises_stopiteration(self):
+        t = tuple[int, str]
+        iter_t = iter(t)
+        next(iter_t)
+        with self.assertRaises(StopIteration):
+            next(iter_t)
+
+    def test_del_iter(self):
+        t = tuple[int, str]
+        iter_x = iter(t)
+        del iter_x
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/Misc/NEWS.d/next/Library/2022-01-30-20-32-40.bpo-43224.zqrQsj.rst b/Misc/NEWS.d/next/Library/2022-01-30-20-32-40.bpo-43224.zqrQsj.rst
@@ -0,0 +1 @@
+Allow unpacking types.GenericAlias objects, e.g. ``*tuple[int, str]``.
diff --git a/Objects/genericaliasobject.c b/Objects/genericaliasobject.c
@@ -5,14 +5,23 @@
 #include "pycore_unionobject.h"   // _Py_union_type_or, _PyGenericAlias_Check
 #include "structmember.h"         // PyMemberDef
 
+#include <stdbool.h>
+
 typedef struct {
     PyObject_HEAD
     PyObject *origin;
     PyObject *args;
     PyObject *parameters;
     PyObject* weakreflist;
+    // Whether we're a starred type, e.g. *tuple[int].
+    bool starred;
 } gaobject;
 
+typedef struct {
+    PyObject_HEAD
+    PyObject *obj;  /* Set to NULL when iterator is exhausted */
+} gaiterobject;
+
 static void
 ga_dealloc(PyObject *self)
 {
@@ -120,6 +129,11 @@ ga_repr(PyObject *self)
     _PyUnicodeWriter writer;
     _PyUnicodeWriter_Init(&writer);
 
+    if (alias->starred) {
+        if (_PyUnicodeWriter_WriteASCIIString(&writer, "*", 1) < 0) {
+            goto error;
+        }
+    }
     if (ga_repr_item(&writer, alias->origin) < 0) {
         goto error;
     }
@@ -603,6 +617,66 @@ static PyNumberMethods ga_as_number = {
         .nb_or = _Py_union_type_or, // Add __or__ function
 };
 
+static PyObject *
+ga_iternext(gaiterobject *gi) {
+    if (gi->obj == NULL) {
+        PyErr_SetNone(PyExc_StopIteration);
+        return NULL;
+    }
+    gaobject *alias = (gaobject *)gi->obj;
+    PyObject *starred_alias = Py_GenericAlias(alias->origin, alias->args);
+    if (starred_alias == NULL) {
+        return NULL;
+    }
+    ((gaobject *)starred_alias)->starred = true;
+    Py_SETREF(gi->obj, NULL);
+    return starred_alias;
+}
+
+static void
+ga_iter_dealloc(gaiterobject *gi) {
+    PyObject_GC_UnTrack(gi);
+    Py_XDECREF(gi->obj);
+    PyObject_GC_Del(gi);
+}
+
+static int
+ga_iter_traverse(gaiterobject *gi, visitproc visit, void *arg)
+{
+    Py_VISIT(gi->obj);
+    return 0;
+}
+
+static int
+ga_iter_clear(PyObject *self) {
+    gaiterobject *gi = (gaiterobject *)self;
+    Py_CLEAR(gi->obj);
+    return 0;
+}
+
+static PyTypeObject Py_GenericAliasIterType = {
+    PyVarObject_HEAD_INIT(&PyType_Type, 0)
+    .tp_name = "generic_alias_iterator",
+    .tp_basicsize = sizeof(gaiterobject),
+    .tp_iter = PyObject_SelfIter,
+    .tp_iternext = (iternextfunc)ga_iternext,
+    .tp_traverse = (traverseproc)ga_iter_traverse,
+    .tp_dealloc = (destructor)ga_iter_dealloc,
+    .tp_clear = (inquiry)ga_iter_clear,
+    .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC,
+};
+
+static PyObject *
+ga_iter(PyObject *self) {
+    gaiterobject *gi = PyObject_GC_New(gaiterobject, &Py_GenericAliasIterType);
+    if (gi == NULL) {
+        return NULL;
+    }
+    gi->obj = Py_NewRef(self);
+    PyObject_GC_Track(gi);
+    return (PyObject *)gi;
+}
+
 // TODO:
 // - argument clinic?
 // - __doc__?
@@ -631,6 +705,7 @@ PyTypeObject Py_GenericAliasType = {
     .tp_new = ga_new,
     .tp_free = PyObject_GC_Del,
     .tp_getset = ga_properties,
+    .tp_iter = (getiterfunc)ga_iter,
 };
 
 PyObject *

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1 @@`
	`1`	+Allow unpacking types.GenericAlias objects, e.g. ``*tuple[int, str]``.