Skip to content

Commit 9915ba2

Browse files
authored
GenericAlias pickle support (#9)
1 parent baf9b32 commit 9915ba2

File tree

2 files changed

+64
-29
lines changed

2 files changed

+64
-29
lines changed

Lib/test/test_genericalias.py

+12
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,18 @@
11
"""Tests for C-implemented GenericAlias."""
22

33
import unittest
4+
import pickle
45
from collections import (
56
defaultdict, deque, OrderedDict, Counter, UserDict, UserList
67
)
78
from collections.abc import *
89
from contextlib import AbstractContextManager, AbstractAsyncContextManager
910
from io import IOBase
1011
from re import Pattern, Match
12+
from types import GenericAlias
1113

14+
from typing import TypeVar
15+
T = TypeVar('T')
1216

1317
class BaseTest(unittest.TestCase):
1418
"""Test basics."""
@@ -172,6 +176,14 @@ class L(list): ...
172176
with self.assertRaises(TypeError):
173177
issubclass(L, list[str])
174178

179+
def test_pickle(self):
180+
alias = GenericAlias(list, T)
181+
s = pickle.dumps(alias)
182+
loaded = pickle.loads(s)
183+
self.assertEqual(alias.__origin__, loaded.__origin__)
184+
self.assertEqual(alias.__args__, loaded.__args__)
185+
self.assertEqual(alias.__parameters__, loaded.__parameters__)
186+
175187

176188
if __name__ == "__main__":
177189
unittest.main()

Objects/descrobject.c

+52-29
Original file line numberDiff line numberDiff line change
@@ -1935,6 +1935,34 @@ tuple_index(PyObject *self, Py_ssize_t len, PyObject *item)
19351935
return -1;
19361936
}
19371937

1938+
// tuple(t for t in args if isinstance(t, TypeVar))
1939+
static PyObject *
1940+
make_parameters(PyObject *args)
1941+
{
1942+
Py_ssize_t len = PyTuple_GET_SIZE(args);
1943+
PyObject *parameters = PyTuple_New(len);
1944+
if (parameters == NULL)
1945+
return NULL;
1946+
Py_ssize_t iparam = 0;
1947+
for (Py_ssize_t iarg = 0; iarg < len; iarg++) {
1948+
PyObject *t = PyTuple_GET_ITEM(args, iarg);
1949+
if (is_typevar(t)) {
1950+
if (tuple_index(parameters, iparam, t) < 0) {
1951+
Py_INCREF(t);
1952+
PyTuple_SET_ITEM(parameters, iparam, t);
1953+
iparam++;
1954+
}
1955+
}
1956+
}
1957+
if (iparam < len) {
1958+
if (_PyTuple_Resize(&parameters, iparam) < 0) {
1959+
Py_XDECREF(parameters);
1960+
return NULL;
1961+
}
1962+
}
1963+
return parameters;
1964+
}
1965+
19381966
static PyObject *
19391967
ga_getitem(PyObject *self, PyObject *item)
19401968
{
@@ -1999,6 +2027,9 @@ static const char* const attr_exceptions[] = {
19992027
"__args__",
20002028
"__parameters__",
20012029
"__mro_entries__",
2030+
"__reduce_ex__", // needed so we don't look up object.__reduce_ex__
2031+
"__reduce__",
2032+
"__setstate__",
20022033
NULL,
20032034
};
20042035

@@ -2077,10 +2108,30 @@ ga_subclasscheck(PyObject *self, PyObject *Py_UNUSED(ignored))
20772108
self);
20782109
}
20792110

2111+
static PyObject *
2112+
ga_reduce(PyObject *self, PyObject *Py_UNUSED(ignored))
2113+
{
2114+
gaobject *alias = (gaobject *)self;
2115+
return Py_BuildValue("O(OO)", Py_TYPE(alias),
2116+
alias->origin, alias->args);
2117+
}
2118+
2119+
static PyObject *
2120+
ga_setstate(PyObject *self, PyObject *state)
2121+
{
2122+
gaobject *alias = (gaobject *)self;
2123+
PyObject *parameters = make_parameters(alias->args);
2124+
Py_INCREF(parameters);
2125+
alias->parameters = parameters;
2126+
Py_RETURN_NONE;
2127+
}
2128+
20802129
static PyMethodDef ga_methods[] = {
20812130
{"__mro_entries__", ga_mro_entries, METH_O},
20822131
{"__instancecheck__", ga_instancecheck, METH_O},
20832132
{"__subclasscheck__", ga_subclasscheck, METH_O},
2133+
{"__reduce__", ga_reduce, METH_NOARGS},
2134+
{"__setstate__", ga_setstate, METH_O},
20842135
{0}
20852136
};
20862137

@@ -2112,7 +2163,7 @@ ga_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
21122163
// - __eq__
21132164
PyTypeObject Py_GenericAliasType = {
21142165
PyVarObject_HEAD_INIT(&PyType_Type, 0)
2115-
.tp_name = "GenericAlias",
2166+
.tp_name = "types.GenericAlias",
21162167
.tp_basicsize = sizeof(gaobject),
21172168
.tp_dealloc = ga_dealloc,
21182169
.tp_repr = ga_repr,
@@ -2129,34 +2180,6 @@ PyTypeObject Py_GenericAliasType = {
21292180
.tp_free = PyObject_GC_Del,
21302181
};
21312182

2132-
// tuple(t for t in args if isinstance(t, TypeVar))
2133-
static PyObject *
2134-
make_parameters(PyObject *args)
2135-
{
2136-
Py_ssize_t len = PyTuple_GET_SIZE(args);
2137-
PyObject *parameters = PyTuple_New(len);
2138-
if (parameters == NULL)
2139-
return NULL;
2140-
Py_ssize_t iparam = 0;
2141-
for (Py_ssize_t iarg = 0; iarg < len; iarg++) {
2142-
PyObject *t = PyTuple_GET_ITEM(args, iarg);
2143-
if (is_typevar(t)) {
2144-
if (tuple_index(parameters, iparam, t) < 0) {
2145-
Py_INCREF(t);
2146-
PyTuple_SET_ITEM(parameters, iparam, t);
2147-
iparam++;
2148-
}
2149-
}
2150-
}
2151-
if (iparam < len) {
2152-
if (_PyTuple_Resize(&parameters, iparam) < 0) {
2153-
Py_XDECREF(parameters);
2154-
return NULL;
2155-
}
2156-
}
2157-
return parameters;
2158-
}
2159-
21602183
PyObject *
21612184
Py_GenericAlias(PyObject *origin, PyObject *args)
21622185
{

0 commit comments

Comments
 (0)