Skip to content

Commit ff77cab

Browse files
pythongh-101277: Add tee data object type to module state
1 parent 343abef commit ff77cab

File tree

2 files changed

+51
-64
lines changed

2 files changed

+51
-64
lines changed

Modules/clinic/itertoolsmodule.c.h

+3-3
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

Modules/itertoolsmodule.c

+48-61
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ typedef struct {
3131
PyTypeObject *starmap_type;
3232
PyTypeObject *takewhile_type;
3333
PyTypeObject *tee_type;
34+
PyTypeObject *teedataobject_type;
3435
PyTypeObject *ziplongest_type;
3536
} itertools_state;
3637

@@ -56,7 +57,7 @@ find_state_by_type(PyTypeObject *tp)
5657
module itertools
5758
class itertools.groupby "groupbyobject *" "clinic_state()->groupby_type"
5859
class itertools._grouper "_grouperobject *" "clinic_state()->_grouper_type"
59-
class itertools.teedataobject "teedataobject *" "&teedataobject_type"
60+
class itertools.teedataobject "teedataobject *" "clinic_state()->teedataobject_type"
6061
class itertools._tee "teeobject *" "clinic_state()->tee_type"
6162
class itertools.batched "batchedobject *" "&batched_type"
6263
class itertools.cycle "cycleobject *" "clinic_state()->cycle_type"
@@ -73,9 +74,8 @@ class itertools.filterfalse "filterfalseobject *" "clinic_state()->filterfalse_t
7374
class itertools.count "countobject *" "clinic_state()->count_type"
7475
class itertools.pairwise "pairwiseobject *" "clinic_state()->pairwise_type"
7576
[clinic start generated code]*/
76-
/*[clinic end generated code: output=da39a3ee5e6b4b0d input=2de6179a0523a80d]*/
77+
/*[clinic end generated code: output=da39a3ee5e6b4b0d input=419423f2d82cf388]*/
7778

78-
static PyTypeObject teedataobject_type;
7979
static PyTypeObject batched_type;
8080

8181
#include "clinic/itertoolsmodule.c.h"
@@ -731,13 +731,13 @@ typedef struct {
731731
} teeobject;
732732

733733
static PyObject *
734-
teedataobject_newinternal(PyObject *it)
734+
teedataobject_newinternal(itertools_state *state, PyObject *it)
735735
{
736-
teedataobject *tdo;
737-
738-
tdo = PyObject_GC_New(teedataobject, &teedataobject_type);
739-
if (tdo == NULL)
736+
teedataobject *tdo = PyObject_GC_New(teedataobject,
737+
state->teedataobject_type);
738+
if (tdo == NULL) {
740739
return NULL;
740+
}
741741

742742
tdo->running = 0;
743743
tdo->numread = 0;
@@ -750,8 +750,11 @@ teedataobject_newinternal(PyObject *it)
750750
static PyObject *
751751
teedataobject_jumplink(teedataobject *tdo)
752752
{
753-
if (tdo->nextlink == NULL)
754-
tdo->nextlink = teedataobject_newinternal(tdo->it);
753+
if (tdo->nextlink == NULL) {
754+
PyTypeObject *tp = Py_TYPE(tdo);
755+
itertools_state *state = find_state_by_type(tp);
756+
tdo->nextlink = teedataobject_newinternal(state, tdo->it);
757+
}
755758
return Py_XNewRef(tdo->nextlink);
756759
}
757760

@@ -787,6 +790,7 @@ teedataobject_traverse(teedataobject *tdo, visitproc visit, void * arg)
787790
{
788791
int i;
789792

793+
Py_VISIT(Py_TYPE(tdo));
790794
Py_VISIT(tdo->it);
791795
for (i = 0; i < tdo->numread; i++)
792796
Py_VISIT(tdo->values[i]);
@@ -795,9 +799,9 @@ teedataobject_traverse(teedataobject *tdo, visitproc visit, void * arg)
795799
}
796800

797801
static void
798-
teedataobject_safe_decref(PyObject *obj)
802+
teedataobject_safe_decref(PyObject *obj, PyTypeObject *tdo_type)
799803
{
800-
while (obj && Py_IS_TYPE(obj, &teedataobject_type) &&
804+
while (obj && Py_IS_TYPE(obj, tdo_type) &&
801805
Py_REFCNT(obj) == 1) {
802806
PyObject *nextlink = ((teedataobject *)obj)->nextlink;
803807
((teedataobject *)obj)->nextlink = NULL;
@@ -817,16 +821,19 @@ teedataobject_clear(teedataobject *tdo)
817821
Py_CLEAR(tdo->values[i]);
818822
tmp = tdo->nextlink;
819823
tdo->nextlink = NULL;
820-
teedataobject_safe_decref(tmp);
824+
itertools_state *state = find_state_by_type(Py_TYPE(tdo));
825+
teedataobject_safe_decref(tmp, state->teedataobject_type);
821826
return 0;
822827
}
823828

824829
static void
825830
teedataobject_dealloc(teedataobject *tdo)
826831
{
832+
PyTypeObject *tp = Py_TYPE(tdo);
827833
PyObject_GC_UnTrack(tdo);
828834
teedataobject_clear(tdo);
829835
PyObject_GC_Del(tdo);
836+
Py_DECREF(tp);
830837
}
831838

832839
static PyObject *
@@ -865,9 +872,10 @@ itertools_teedataobject_impl(PyTypeObject *type, PyObject *it,
865872
teedataobject *tdo;
866873
Py_ssize_t i, len;
867874

868-
assert(type == &teedataobject_type);
875+
itertools_state *state = find_state_by_type(type);
876+
assert(type == state->teedataobject_type);
869877

870-
tdo = (teedataobject *)teedataobject_newinternal(it);
878+
tdo = (teedataobject *)teedataobject_newinternal(state, it);
871879
if (!tdo)
872880
return NULL;
873881

@@ -883,7 +891,7 @@ itertools_teedataobject_impl(PyTypeObject *type, PyObject *it,
883891

884892
if (len == LINKCELLS) {
885893
if (next != Py_None) {
886-
if (!Py_IS_TYPE(next, &teedataobject_type))
894+
if (!Py_IS_TYPE(next, state->teedataobject_type))
887895
goto err;
888896
assert(tdo->nextlink == NULL);
889897
tdo->nextlink = Py_NewRef(next);
@@ -906,47 +914,24 @@ static PyMethodDef teedataobject_methods[] = {
906914
{NULL, NULL} /* sentinel */
907915
};
908916

909-
static PyTypeObject teedataobject_type = {
910-
PyVarObject_HEAD_INIT(0, 0) /* Must fill in type value later */
911-
"itertools._tee_dataobject", /* tp_name */
912-
sizeof(teedataobject), /* tp_basicsize */
913-
0, /* tp_itemsize */
914-
/* methods */
915-
(destructor)teedataobject_dealloc, /* tp_dealloc */
916-
0, /* tp_vectorcall_offset */
917-
0, /* tp_getattr */
918-
0, /* tp_setattr */
919-
0, /* tp_as_async */
920-
0, /* tp_repr */
921-
0, /* tp_as_number */
922-
0, /* tp_as_sequence */
923-
0, /* tp_as_mapping */
924-
0, /* tp_hash */
925-
0, /* tp_call */
926-
0, /* tp_str */
927-
PyObject_GenericGetAttr, /* tp_getattro */
928-
0, /* tp_setattro */
929-
0, /* tp_as_buffer */
930-
Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC, /* tp_flags */
931-
itertools_teedataobject__doc__, /* tp_doc */
932-
(traverseproc)teedataobject_traverse, /* tp_traverse */
933-
(inquiry)teedataobject_clear, /* tp_clear */
934-
0, /* tp_richcompare */
935-
0, /* tp_weaklistoffset */
936-
0, /* tp_iter */
937-
0, /* tp_iternext */
938-
teedataobject_methods, /* tp_methods */
939-
0, /* tp_members */
940-
0, /* tp_getset */
941-
0, /* tp_base */
942-
0, /* tp_dict */
943-
0, /* tp_descr_get */
944-
0, /* tp_descr_set */
945-
0, /* tp_dictoffset */
946-
0, /* tp_init */
947-
0, /* tp_alloc */
948-
itertools_teedataobject, /* tp_new */
949-
PyObject_GC_Del, /* tp_free */
917+
static PyType_Slot teedataobject_slots[] = {
918+
{Py_tp_dealloc, teedataobject_dealloc},
919+
{Py_tp_getattro, PyObject_GenericGetAttr},
920+
{Py_tp_doc, (void *)itertools_teedataobject__doc__},
921+
{Py_tp_traverse, teedataobject_traverse},
922+
{Py_tp_clear, teedataobject_clear},
923+
{Py_tp_methods, teedataobject_methods},
924+
{Py_tp_new, itertools_teedataobject},
925+
{Py_tp_free, PyObject_GC_Del},
926+
{0, NULL},
927+
};
928+
929+
static PyType_Spec teedataobject_spec = {
930+
.name = "itertools._tee_dataobject",
931+
.basicsize = sizeof(teedataobject),
932+
.flags = (Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HAVE_GC |
933+
Py_TPFLAGS_IMMUTABLETYPE),
934+
.slots = teedataobject_slots,
950935
};
951936

952937

@@ -1007,7 +992,7 @@ tee_fromiterable(itertools_state *state, PyObject *iterable)
1007992
goto done;
1008993
}
1009994

1010-
PyObject *dataobj = teedataobject_newinternal(it);
995+
PyObject *dataobj = teedataobject_newinternal(state, it);
1011996
if (!dataobj) {
1012997
to = NULL;
1013998
goto done;
@@ -1076,7 +1061,9 @@ tee_setstate(teeobject *to, PyObject *state)
10761061
PyErr_SetString(PyExc_TypeError, "state is not a tuple");
10771062
return NULL;
10781063
}
1079-
if (!PyArg_ParseTuple(state, "O!i", &teedataobject_type, &tdo, &index)) {
1064+
itertools_state *m_state = find_state_by_type(Py_TYPE(to));
1065+
PyTypeObject *tdo_type = m_state->teedataobject_type;
1066+
if (!PyArg_ParseTuple(state, "O!i", tdo_type, &tdo, &index)) {
10801067
return NULL;
10811068
}
10821069
if (index < 0 || index > LINKCELLS) {
@@ -4696,14 +4683,14 @@ itertoolsmodule_exec(PyObject *mod)
46964683
ADD_TYPE(mod, state->starmap_type, &starmap_spec);
46974684
ADD_TYPE(mod, state->takewhile_type, &takewhile_spec);
46984685
ADD_TYPE(mod, state->tee_type, &tee_spec);
4686+
ADD_TYPE(mod, state->teedataobject_type, &teedataobject_spec);
46994687
ADD_TYPE(mod, state->ziplongest_type, &ziplongest_spec);
47004688

47014689
PyTypeObject *typelist[] = {
47024690
&batched_type,
4703-
&teedataobject_type
47044691
};
47054692

4706-
Py_SET_TYPE(&teedataobject_type, &PyType_Type);
4693+
Py_SET_TYPE(state->teedataobject_type, &PyType_Type);
47074694

47084695
for (size_t i = 0; i < Py_ARRAY_LENGTH(typelist); i++) {
47094696
if (PyModule_AddType(mod, typelist[i]) < 0) {

0 commit comments

Comments
 (0)