From 9915ba2bea0c1eb6b2f5838521dbebcaf8c5c19b Mon Sep 17 00:00:00 2001
From: Ethan Smith <ethan@ethanhs.me>
Date: Tue, 4 Feb 2020 20:36:28 -0800
Subject: [PATCH] GenericAlias pickle support (#9)

---
 Lib/test/test_genericalias.py | 12 ++++++
 Objects/descrobject.c         | 81 ++++++++++++++++++++++-------------
 2 files changed, 64 insertions(+), 29 deletions(-)

diff --git a/Lib/test/test_genericalias.py b/Lib/test/test_genericalias.py
index 94c32dc976aa66..0b6413d3010686 100644
--- a/Lib/test/test_genericalias.py
+++ b/Lib/test/test_genericalias.py
@@ -1,6 +1,7 @@
 """Tests for C-implemented GenericAlias."""
 
 import unittest
+import pickle
 from collections import (
     defaultdict, deque, OrderedDict, Counter, UserDict, UserList
 )
@@ -8,7 +9,10 @@
 from contextlib import AbstractContextManager, AbstractAsyncContextManager
 from io import IOBase
 from re import Pattern, Match
+from types import GenericAlias
 
+from typing import TypeVar
+T = TypeVar('T')
 
 class BaseTest(unittest.TestCase):
     """Test basics."""
@@ -172,6 +176,14 @@ class L(list): ...
         with self.assertRaises(TypeError):
             issubclass(L, list[str])
 
+    def test_pickle(self):
+        alias = GenericAlias(list, T)
+        s = pickle.dumps(alias)
+        loaded = pickle.loads(s)
+        self.assertEqual(alias.__origin__, loaded.__origin__)
+        self.assertEqual(alias.__args__, loaded.__args__)
+        self.assertEqual(alias.__parameters__, loaded.__parameters__)
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/Objects/descrobject.c b/Objects/descrobject.c
index 2e34f9fe0da366..747ce75d471c0c 100644
--- a/Objects/descrobject.c
+++ b/Objects/descrobject.c
@@ -1935,6 +1935,34 @@ tuple_index(PyObject *self, Py_ssize_t len, PyObject *item)
     return -1;
 }
 
+// tuple(t for t in args if isinstance(t, TypeVar))
+static PyObject *
+make_parameters(PyObject *args)
+{
+    Py_ssize_t len = PyTuple_GET_SIZE(args);
+    PyObject *parameters = PyTuple_New(len);
+    if (parameters == NULL)
+        return NULL;
+    Py_ssize_t iparam = 0;
+    for (Py_ssize_t iarg = 0; iarg < len; iarg++) {
+        PyObject *t = PyTuple_GET_ITEM(args, iarg);
+        if (is_typevar(t)) {
+            if (tuple_index(parameters, iparam, t) < 0) {
+                Py_INCREF(t);
+                PyTuple_SET_ITEM(parameters, iparam, t);
+                iparam++;
+            }
+        }
+    }
+    if (iparam < len) {
+        if (_PyTuple_Resize(&parameters, iparam) < 0) {
+            Py_XDECREF(parameters);
+            return NULL;
+        }
+    }
+    return parameters;
+}
+
 static PyObject *
 ga_getitem(PyObject *self, PyObject *item)
 {
@@ -1999,6 +2027,9 @@ static const char* const attr_exceptions[] = {
     "__args__",
     "__parameters__",
     "__mro_entries__",
+    "__reduce_ex__",  // needed so we don't look up object.__reduce_ex__
+    "__reduce__",
+    "__setstate__",
     NULL,
 };
 
@@ -2077,10 +2108,30 @@ ga_subclasscheck(PyObject *self, PyObject *Py_UNUSED(ignored))
                         self);
 }
 
+static PyObject *
+ga_reduce(PyObject *self, PyObject *Py_UNUSED(ignored))
+{
+    gaobject *alias = (gaobject *)self;
+    return Py_BuildValue("O(OO)", Py_TYPE(alias),
+                         alias->origin, alias->args);
+}
+
+static PyObject *
+ga_setstate(PyObject *self, PyObject *state)
+{
+    gaobject *alias = (gaobject *)self;
+    PyObject *parameters = make_parameters(alias->args);
+    Py_INCREF(parameters);
+    alias->parameters = parameters;
+    Py_RETURN_NONE;
+}
+
 static PyMethodDef ga_methods[] = {
     {"__mro_entries__", ga_mro_entries, METH_O},
     {"__instancecheck__", ga_instancecheck, METH_O},
     {"__subclasscheck__", ga_subclasscheck, METH_O},
+    {"__reduce__", ga_reduce, METH_NOARGS},
+    {"__setstate__", ga_setstate, METH_O},
     {0}
 };
 
@@ -2112,7 +2163,7 @@ ga_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
 // - __eq__
 PyTypeObject Py_GenericAliasType = {
     PyVarObject_HEAD_INIT(&PyType_Type, 0)
-    .tp_name = "GenericAlias",
+    .tp_name = "types.GenericAlias",
     .tp_basicsize = sizeof(gaobject),
     .tp_dealloc = ga_dealloc,
     .tp_repr = ga_repr,
@@ -2129,34 +2180,6 @@ PyTypeObject Py_GenericAliasType = {
     .tp_free = PyObject_GC_Del,
 };
 
-// tuple(t for t in args if isinstance(t, TypeVar))
-static PyObject *
-make_parameters(PyObject *args)
-{
-    Py_ssize_t len = PyTuple_GET_SIZE(args);
-    PyObject *parameters = PyTuple_New(len);
-    if (parameters == NULL)
-        return NULL;
-    Py_ssize_t iparam = 0;
-    for (Py_ssize_t iarg = 0; iarg < len; iarg++) {
-        PyObject *t = PyTuple_GET_ITEM(args, iarg);
-        if (is_typevar(t)) {
-            if (tuple_index(parameters, iparam, t) < 0) {
-                Py_INCREF(t);
-                PyTuple_SET_ITEM(parameters, iparam, t);
-                iparam++;
-            }
-        }
-    }
-    if (iparam < len) {
-        if (_PyTuple_Resize(&parameters, iparam) < 0) {
-            Py_XDECREF(parameters);
-            return NULL;
-        }
-    }
-    return parameters;
-}
-
 PyObject *
 Py_GenericAlias(PyObject *origin, PyObject *args)
 {