Skip to content

Commit b7c21be

Browse files
erlend-aaslanderlend-aasland
committed
pythongh-101277: Add tee data object type to module state
1 parent 75bb5da commit b7c21be

File tree

2 files changed

+50
-63
lines changed

2 files changed

+50
-63
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: 48 additions & 61 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@ typedef struct {
3232
PyTypeObject *starmap_type;
3333
PyTypeObject *takewhile_type;
3434
PyTypeObject *tee_type;
35+
PyTypeObject *teedataobject_type;
3536
PyTypeObject *ziplongest_type;
3637
} itertools_state;
3738

@@ -66,7 +67,7 @@ find_state_by_type(PyTypeObject *tp)
6667
module itertools
6768
class itertools.groupby "groupbyobject *" "clinic_state()->groupby_type"
6869
class itertools._grouper "_grouperobject *" "clinic_state()->_grouper_type"
69-
class itertools.teedataobject "teedataobject *" "&teedataobject_type"
70+
class itertools.teedataobject "teedataobject *" "clinic_state()->teedataobject_type"
7071
class itertools._tee "teeobject *" "clinic_state()->tee_type"
7172
class itertools.batched "batchedobject *" "&batched_type"
7273
class itertools.cycle "cycleobject *" "clinic_state()->cycle_type"
@@ -83,9 +84,8 @@ class itertools.filterfalse "filterfalseobject *" "clinic_state()->filterfalse_t
8384
class itertools.count "countobject *" "clinic_state()->count_type"
8485
class itertools.pairwise "pairwiseobject *" "clinic_state()->pairwise_type"
8586
[clinic start generated code]*/
86-
/*[clinic end generated code: output=da39a3ee5e6b4b0d input=2de6179a0523a80d]*/
87+
/*[clinic end generated code: output=da39a3ee5e6b4b0d input=419423f2d82cf388]*/
8788

88-
static PyTypeObject teedataobject_type;
8989
static PyTypeObject batched_type;
9090

9191
#define clinic_state_by_cls() (get_module_state_by_cls(base_tp))
@@ -744,13 +744,13 @@ typedef struct {
744744
} teeobject;
745745

746746
static PyObject *
747-
teedataobject_newinternal(PyObject *it)
747+
teedataobject_newinternal(itertools_state *state, PyObject *it)
748748
{
749-
teedataobject *tdo;
750-
751-
tdo = PyObject_GC_New(teedataobject, &teedataobject_type);
752-
if (tdo == NULL)
749+
teedataobject *tdo = PyObject_GC_New(teedataobject,
750+
state->teedataobject_type);
751+
if (tdo == NULL) {
753752
return NULL;
753+
}
754754

755755
tdo->running = 0;
756756
tdo->numread = 0;
@@ -763,8 +763,11 @@ teedataobject_newinternal(PyObject *it)
763763
static PyObject *
764764
teedataobject_jumplink(teedataobject *tdo)
765765
{
766-
if (tdo->nextlink == NULL)
767-
tdo->nextlink = teedataobject_newinternal(tdo->it);
766+
if (tdo->nextlink == NULL) {
767+
PyTypeObject *tp = Py_TYPE(tdo);
768+
itertools_state *state = find_state_by_type(tp);
769+
tdo->nextlink = teedataobject_newinternal(state, tdo->it);
770+
}
768771
return Py_XNewRef(tdo->nextlink);
769772
}
770773

@@ -800,6 +803,7 @@ teedataobject_traverse(teedataobject *tdo, visitproc visit, void * arg)
800803
{
801804
int i;
802805

806+
Py_VISIT(Py_TYPE(tdo));
803807
Py_VISIT(tdo->it);
804808
for (i = 0; i < tdo->numread; i++)
805809
Py_VISIT(tdo->values[i]);
@@ -808,9 +812,9 @@ teedataobject_traverse(teedataobject *tdo, visitproc visit, void * arg)
808812
}
809813

810814
static void
811-
teedataobject_safe_decref(PyObject *obj)
815+
teedataobject_safe_decref(PyObject *obj, PyTypeObject *tdo_type)
812816
{
813-
while (obj && Py_IS_TYPE(obj, &teedataobject_type) &&
817+
while (obj && Py_IS_TYPE(obj, tdo_type) &&
814818
Py_REFCNT(obj) == 1) {
815819
PyObject *nextlink = ((teedataobject *)obj)->nextlink;
816820
((teedataobject *)obj)->nextlink = NULL;
@@ -830,16 +834,19 @@ teedataobject_clear(teedataobject *tdo)
830834
Py_CLEAR(tdo->values[i]);
831835
tmp = tdo->nextlink;
832836
tdo->nextlink = NULL;
833-
teedataobject_safe_decref(tmp);
837+
itertools_state *state = find_state_by_type(Py_TYPE(tdo));
838+
teedataobject_safe_decref(tmp, state->teedataobject_type);
834839
return 0;
835840
}
836841

837842
static void
838843
teedataobject_dealloc(teedataobject *tdo)
839844
{
845+
PyTypeObject *tp = Py_TYPE(tdo);
840846
PyObject_GC_UnTrack(tdo);
841847
teedataobject_clear(tdo);
842848
PyObject_GC_Del(tdo);
849+
Py_DECREF(tp);
843850
}
844851

845852
static PyObject *
@@ -878,9 +885,10 @@ itertools_teedataobject_impl(PyTypeObject *type, PyObject *it,
878885
teedataobject *tdo;
879886
Py_ssize_t i, len;
880887

881-
assert(type == &teedataobject_type);
888+
itertools_state *state = find_state_by_type(type);
889+
assert(type == state->teedataobject_type);
882890

883-
tdo = (teedataobject *)teedataobject_newinternal(it);
891+
tdo = (teedataobject *)teedataobject_newinternal(state, it);
884892
if (!tdo)
885893
return NULL;
886894

@@ -896,7 +904,7 @@ itertools_teedataobject_impl(PyTypeObject *type, PyObject *it,
896904

897905
if (len == LINKCELLS) {
898906
if (next != Py_None) {
899-
if (!Py_IS_TYPE(next, &teedataobject_type))
907+
if (!Py_IS_TYPE(next, state->teedataobject_type))
900908
goto err;
901909
assert(tdo->nextlink == NULL);
902910
tdo->nextlink = Py_NewRef(next);
@@ -919,47 +927,24 @@ static PyMethodDef teedataobject_methods[] = {
919927
{NULL, NULL} /* sentinel */
920928
};
921929

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

965950

@@ -1020,7 +1005,7 @@ tee_fromiterable(itertools_state *state, PyObject *iterable)
10201005
goto done;
10211006
}
10221007

1023-
PyObject *dataobj = teedataobject_newinternal(it);
1008+
PyObject *dataobj = teedataobject_newinternal(state, it);
10241009
if (!dataobj) {
10251010
to = NULL;
10261011
goto done;
@@ -1089,7 +1074,9 @@ tee_setstate(teeobject *to, PyObject *state)
10891074
PyErr_SetString(PyExc_TypeError, "state is not a tuple");
10901075
return NULL;
10911076
}
1092-
if (!PyArg_ParseTuple(state, "O!i", &teedataobject_type, &tdo, &index)) {
1077+
itertools_state *m_state = find_state_by_type(Py_TYPE(to));
1078+
PyTypeObject *tdo_type = m_state->teedataobject_type;
1079+
if (!PyArg_ParseTuple(state, "O!i", tdo_type, &tdo, &index)) {
10931080
return NULL;
10941081
}
10951082
if (index < 0 || index > LINKCELLS) {
@@ -4712,14 +4699,14 @@ itertoolsmodule_exec(PyObject *mod)
47124699
ADD_TYPE(mod, state->starmap_type, &starmap_spec);
47134700
ADD_TYPE(mod, state->takewhile_type, &takewhile_spec);
47144701
ADD_TYPE(mod, state->tee_type, &tee_spec);
4702+
ADD_TYPE(mod, state->teedataobject_type, &teedataobject_spec);
47154703
ADD_TYPE(mod, state->ziplongest_type, &ziplongest_spec);
47164704

47174705
PyTypeObject *typelist[] = {
47184706
&batched_type,
4719-
&teedataobject_type
47204707
};
47214708

4722-
Py_SET_TYPE(&teedataobject_type, &PyType_Type);
4709+
Py_SET_TYPE(state->teedataobject_type, &PyType_Type);
47234710

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

0 commit comments

Comments
 (0)