Skip to content

Commit e574de1

Browse files
pythongh-101277: Add tee type to module state
1 parent d18169b commit e574de1

File tree

2 files changed

+47
-56
lines changed

2 files changed

+47
-56
lines changed

Modules/clinic/itertoolsmodule.c.h

Lines changed: 2 additions & 2 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

Modules/itertoolsmodule.c

Lines changed: 45 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
#include "pycore_moduleobject.h" // _PyModule_GetState()
66
#include "pycore_object.h" // _PyObject_GC_TRACK()
77
#include "pycore_tuple.h" // _PyTuple_ITEMS()
8+
#include "structmember.h" // PyMemberDef
89
#include <stddef.h> // offsetof()
910

1011
/* Itertools module written and maintained
@@ -30,6 +31,7 @@ typedef struct {
3031
PyTypeObject *repeat_type;
3132
PyTypeObject *starmap_type;
3233
PyTypeObject *takewhile_type;
34+
PyTypeObject *tee_type;
3335
PyTypeObject *ziplongest_type;
3436
} itertools_state;
3537

@@ -64,7 +66,7 @@ module itertools
6466
class itertools.groupby "groupbyobject *" "clinic_state()->groupby_type"
6567
class itertools._grouper "_grouperobject *" "clinic_state()->_grouper_type"
6668
class itertools.teedataobject "teedataobject *" "&teedataobject_type"
67-
class itertools._tee "teeobject *" "&tee_type"
69+
class itertools._tee "teeobject *" "clinic_state()->tee_type"
6870
class itertools.batched "batchedobject *" "&batched_type"
6971
class itertools.cycle "cycleobject *" "clinic_state()->cycle_type"
7072
class itertools.dropwhile "dropwhileobject *" "clinic_state()->dropwhile_type"
@@ -80,10 +82,9 @@ class itertools.filterfalse "filterfalseobject *" "clinic_state()->filterfalse_t
8082
class itertools.count "countobject *" "clinic_state()->count_type"
8183
class itertools.pairwise "pairwiseobject *" "clinic_state()->pairwise_type"
8284
[clinic start generated code]*/
83-
/*[clinic end generated code: output=da39a3ee5e6b4b0d input=102319065d3e090b]*/
85+
/*[clinic end generated code: output=da39a3ee5e6b4b0d input=2de6179a0523a80d]*/
8486

8587
static PyTypeObject teedataobject_type;
86-
static PyTypeObject tee_type;
8788
static PyTypeObject batched_type;
8889

8990
#define clinic_state() (find_state_by_type(type))
@@ -984,18 +985,18 @@ tee_next(teeobject *to)
984985
static int
985986
tee_traverse(teeobject *to, visitproc visit, void *arg)
986987
{
988+
Py_VISIT(Py_TYPE(to));
987989
Py_VISIT((PyObject *)to->dataobj);
988990
return 0;
989991
}
990992

991993
static PyObject *
992994
tee_copy(teeobject *to, PyObject *Py_UNUSED(ignored))
993995
{
994-
teeobject *newto;
995-
996-
newto = PyObject_GC_New(teeobject, &tee_type);
997-
if (newto == NULL)
996+
teeobject *newto = PyObject_GC_New(teeobject, Py_TYPE(to));
997+
if (newto == NULL) {
998998
return NULL;
999+
}
9991000
newto->dataobj = (teedataobject*)Py_NewRef(to->dataobj);
10001001
newto->index = to->index;
10011002
newto->weakreflist = NULL;
@@ -1006,15 +1007,15 @@ tee_copy(teeobject *to, PyObject *Py_UNUSED(ignored))
10061007
PyDoc_STRVAR(teecopy_doc, "Returns an independent iterator.");
10071008

10081009
static PyObject *
1009-
tee_fromiterable(PyObject *iterable)
1010+
tee_fromiterable(itertools_state *state, PyObject *iterable)
10101011
{
10111012
teeobject *to;
10121013
PyObject *it;
10131014

10141015
it = PyObject_GetIter(iterable);
10151016
if (it == NULL)
10161017
return NULL;
1017-
if (PyObject_TypeCheck(it, &tee_type)) {
1018+
if (PyObject_TypeCheck(it, state->tee_type)) {
10181019
to = (teeobject *)tee_copy((teeobject *)it, NULL);
10191020
goto done;
10201021
}
@@ -1024,7 +1025,7 @@ tee_fromiterable(PyObject *iterable)
10241025
to = NULL;
10251026
goto done;
10261027
}
1027-
to = PyObject_GC_New(teeobject, &tee_type);
1028+
to = PyObject_GC_New(teeobject, state->tee_type);
10281029
if (to == NULL) {
10291030
Py_DECREF(dataobj);
10301031
goto done;
@@ -1050,7 +1051,8 @@ static PyObject *
10501051
itertools__tee_impl(PyTypeObject *type, PyObject *iterable)
10511052
/*[clinic end generated code: output=b02d3fd26c810c3f input=adc0779d2afe37a2]*/
10521053
{
1053-
return tee_fromiterable(iterable);
1054+
itertools_state *state = find_state_by_type(type);
1055+
return tee_fromiterable(state, iterable);
10541056
}
10551057

10561058
static int
@@ -1065,9 +1067,11 @@ tee_clear(teeobject *to)
10651067
static void
10661068
tee_dealloc(teeobject *to)
10671069
{
1070+
PyTypeObject *tp = Py_TYPE(to);
10681071
PyObject_GC_UnTrack(to);
10691072
tee_clear(to);
10701073
PyObject_GC_Del(to);
1074+
Py_DECREF(tp);
10711075
}
10721076

10731077
static PyObject *
@@ -1105,47 +1109,31 @@ static PyMethodDef tee_methods[] = {
11051109
{NULL, NULL} /* sentinel */
11061110
};
11071111

1108-
static PyTypeObject tee_type = {
1109-
PyVarObject_HEAD_INIT(NULL, 0)
1110-
"itertools._tee", /* tp_name */
1111-
sizeof(teeobject), /* tp_basicsize */
1112-
0, /* tp_itemsize */
1113-
/* methods */
1114-
(destructor)tee_dealloc, /* tp_dealloc */
1115-
0, /* tp_vectorcall_offset */
1116-
0, /* tp_getattr */
1117-
0, /* tp_setattr */
1118-
0, /* tp_as_async */
1119-
0, /* tp_repr */
1120-
0, /* tp_as_number */
1121-
0, /* tp_as_sequence */
1122-
0, /* tp_as_mapping */
1123-
0, /* tp_hash */
1124-
0, /* tp_call */
1125-
0, /* tp_str */
1126-
0, /* tp_getattro */
1127-
0, /* tp_setattro */
1128-
0, /* tp_as_buffer */
1129-
Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
1130-
itertools__tee__doc__, /* tp_doc */
1131-
(traverseproc)tee_traverse, /* tp_traverse */
1132-
(inquiry)tee_clear, /* tp_clear */
1133-
0, /* tp_richcompare */
1134-
offsetof(teeobject, weakreflist), /* tp_weaklistoffset */
1135-
PyObject_SelfIter, /* tp_iter */
1136-
(iternextfunc)tee_next, /* tp_iternext */
1137-
tee_methods, /* tp_methods */
1138-
0, /* tp_members */
1139-
0, /* tp_getset */
1140-
0, /* tp_base */
1141-
0, /* tp_dict */
1142-
0, /* tp_descr_get */
1143-
0, /* tp_descr_set */
1144-
0, /* tp_dictoffset */
1145-
0, /* tp_init */
1146-
0, /* tp_alloc */
1147-
itertools__tee, /* tp_new */
1148-
PyObject_GC_Del, /* tp_free */
1112+
static PyMemberDef tee_members[] = {
1113+
{"__weaklistoffset__", T_PYSSIZET, offsetof(teeobject, weakreflist), READONLY},
1114+
{NULL},
1115+
};
1116+
1117+
static PyType_Slot tee_slots[] = {
1118+
{Py_tp_dealloc, tee_dealloc},
1119+
{Py_tp_doc, (void *)itertools__tee__doc__},
1120+
{Py_tp_traverse, tee_traverse},
1121+
{Py_tp_clear, tee_clear},
1122+
{Py_tp_iter, PyObject_SelfIter},
1123+
{Py_tp_iternext, tee_next},
1124+
{Py_tp_methods, tee_methods},
1125+
{Py_tp_members, tee_members},
1126+
{Py_tp_new, itertools__tee},
1127+
{Py_tp_free, PyObject_GC_Del},
1128+
{0, NULL},
1129+
};
1130+
1131+
static PyType_Spec tee_spec = {
1132+
.name = "itertools._tee",
1133+
.basicsize = sizeof(teeobject),
1134+
.flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
1135+
Py_TPFLAGS_IMMUTABLETYPE),
1136+
.slots = tee_slots,
11491137
};
11501138

11511139
/*[clinic input]
@@ -1187,7 +1175,8 @@ itertools_tee_impl(PyObject *module, PyObject *iterable, Py_ssize_t n)
11871175
copyable = it;
11881176
}
11891177
else {
1190-
copyable = tee_fromiterable(it);
1178+
itertools_state *state = get_module_state(module);
1179+
copyable = tee_fromiterable(state, it);
11911180
Py_DECREF(it);
11921181
if (copyable == NULL) {
11931182
Py_DECREF(result);
@@ -4651,6 +4640,7 @@ itertoolsmodule_traverse(PyObject *mod, visitproc visit, void *arg)
46514640
Py_VISIT(state->repeat_type);
46524641
Py_VISIT(state->starmap_type);
46534642
Py_VISIT(state->takewhile_type);
4643+
Py_VISIT(state->tee_type);
46544644
Py_VISIT(state->ziplongest_type);
46554645
return 0;
46564646
}
@@ -4677,6 +4667,7 @@ itertoolsmodule_clear(PyObject *mod)
46774667
Py_CLEAR(state->repeat_type);
46784668
Py_CLEAR(state->starmap_type);
46794669
Py_CLEAR(state->takewhile_type);
4670+
Py_CLEAR(state->tee_type);
46804671
Py_CLEAR(state->ziplongest_type);
46814672
return 0;
46824673
}
@@ -4720,11 +4711,11 @@ itertoolsmodule_exec(PyObject *mod)
47204711
ADD_TYPE(mod, state->repeat_type, &repeat_spec);
47214712
ADD_TYPE(mod, state->starmap_type, &starmap_spec);
47224713
ADD_TYPE(mod, state->takewhile_type, &takewhile_spec);
4714+
ADD_TYPE(mod, state->tee_type, &tee_spec);
47234715
ADD_TYPE(mod, state->ziplongest_type, &ziplongest_spec);
47244716

47254717
PyTypeObject *typelist[] = {
47264718
&batched_type,
4727-
&tee_type,
47284719
&teedataobject_type
47294720
};
47304721

0 commit comments

Comments
 (0)