Skip to content

Commit a9d57fc

Browse files
emmatypingemmatyping
committed
Make objects GC tracked
1 parent c786c3c commit a9d57fc

File tree

3 files changed

+86
-19
lines changed

3 files changed

+86
-19
lines changed

‎Modules/_zstd/compressor.c‎

Lines changed: 28 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -299,7 +299,7 @@ static PyObject *
299299
_zstd_ZstdCompressor_new(PyTypeObject *type, PyObject *Py_UNUSED(args), PyObject *Py_UNUSED(kwargs))
300300
{
301301
ZstdCompressor *self;
302-
self = (ZstdCompressor*)type->tp_alloc(type, 0);
302+
self = PyObject_GC_New(ZstdCompressor, type);
303303
if (self == NULL) {
304304
goto error;
305305
}
@@ -322,22 +322,25 @@ _zstd_ZstdCompressor_new(PyTypeObject *type, PyObject *Py_UNUSED(args), PyObject
322322
return (PyObject*)self;
323323

324324
error:
325-
Py_XDECREF(self);
325+
PyObject_GC_UnTrack(self);
326326
return NULL;
327327
}
328328

329329
static void
330330
ZstdCompressor_dealloc(PyObject *ob)
331331
{
332332
ZstdCompressor *self = ZstdCompressor_CAST(ob);
333+
334+
PyObject_GC_UnTrack(self);
335+
333336
/* Free compression context */
334337
ZSTD_freeCCtx(self->cctx);
335338

336339
/* Py_XDECREF the dict after free the compression context */
337-
Py_XDECREF(self->dict);
340+
Py_CLEAR(self->dict);
338341

339342
PyTypeObject *tp = Py_TYPE(self);
340-
tp->tp_free((PyObject*)self);
343+
PyObject_GC_Del(ob);
341344
Py_DECREF(tp);
342345
}
343346

@@ -398,6 +401,8 @@ _zstd_ZstdCompressor___init___impl(ZstdCompressor *self, PyObject *level,
398401
self->dict = zstd_dict;
399402
}
400403

404+
// We can only start tracking self with the GC once self->dict is set.
405+
PyObject_GC_Track(self);
401406
return 0;
402407
}
403408

@@ -660,19 +665,37 @@ static PyMemberDef ZstdCompressor_members[] = {
660665
{0}
661666
};
662667

668+
static int
669+
ZstdCompressor_traverse(PyObject *ob, visitproc visit, void *arg)
670+
{
671+
ZstdCompressor *self = ZstdCompressor_CAST(ob);
672+
Py_VISIT(self->dict);
673+
return 0;
674+
}
675+
676+
static int
677+
ZstdCompressor_clear(PyObject *ob)
678+
{
679+
ZstdCompressor *self = ZstdCompressor_CAST(ob);
680+
Py_CLEAR(self->dict);
681+
return 0;
682+
}
683+
663684
static PyType_Slot zstdcompressor_slots[] = {
664685
{Py_tp_new, _zstd_ZstdCompressor_new},
665686
{Py_tp_dealloc, ZstdCompressor_dealloc},
666687
{Py_tp_init, _zstd_ZstdCompressor___init__},
667688
{Py_tp_methods, ZstdCompressor_methods},
668689
{Py_tp_members, ZstdCompressor_members},
669690
{Py_tp_doc, (char*)_zstd_ZstdCompressor___init____doc__},
691+
{Py_tp_traverse, ZstdCompressor_traverse},
692+
{Py_tp_clear, ZstdCompressor_clear},
670693
{0}
671694
};
672695

673696
PyType_Spec zstdcompressor_type_spec = {
674697
.name = "_zstd.ZstdCompressor",
675698
.basicsize = sizeof(ZstdCompressor),
676-
.flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE,
699+
.flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC,
677700
.slots = zstdcompressor_slots,
678701
};

‎Modules/_zstd/decompressor.c‎

Lines changed: 31 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -617,7 +617,7 @@ static PyObject *
617617
_zstd_ZstdDecompressor_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
618618
{
619619
ZstdDecompressor *self;
620-
self = (ZstdDecompressor*)type->tp_alloc(type, 0);
620+
self = PyObject_GC_New(ZstdDecompressor, type);
621621
if (self == NULL) {
622622
goto error;
623623
}
@@ -642,28 +642,31 @@ _zstd_ZstdDecompressor_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
642642
return (PyObject*)self;
643643

644644
error:
645-
Py_XDECREF(self);
645+
PyObject_GC_UnTrack(self);
646646
return NULL;
647647
}
648648

649649
static void
650650
ZstdDecompressor_dealloc(PyObject *ob)
651651
{
652652
ZstdDecompressor *self = ZstdDecompressor_CAST(ob);
653+
654+
PyObject_GC_UnTrack(self);
655+
653656
/* Free decompression context */
654657
ZSTD_freeDCtx(self->dctx);
655658

656-
/* Py_XDECREF the dict after free decompression context */
657-
Py_XDECREF(self->dict);
659+
/* Py_CLEAR the dict after free decompression context */
660+
Py_CLEAR(self->dict);
658661

659662
/* Free unconsumed input data buffer */
660663
PyMem_Free(self->input_buffer);
661664

662665
/* Free unused data */
663-
Py_XDECREF(self->unused_data);
666+
Py_CLEAR(self->unused_data);
664667

665668
PyTypeObject *tp = Py_TYPE(self);
666-
tp->tp_free((PyObject*)self);
669+
PyObject_GC_Del(ob);
667670
Py_DECREF(tp);
668671
}
669672

@@ -711,6 +714,8 @@ _zstd_ZstdDecompressor___init___impl(ZstdDecompressor *self,
711714
}
712715
}
713716

717+
// We can only start tracking self with the GC once self->dict is set.
718+
PyObject_GC_Track(self);
714719
return 0;
715720
}
716721

@@ -837,6 +842,23 @@ static PyGetSetDef ZstdDecompressor_getset[] = {
837842
{0}
838843
};
839844

845+
static int
846+
ZstdDecompressor_traverse(PyObject *ob, visitproc visit, void *arg)
847+
{
848+
ZstdDecompressor *self = ZstdDecompressor_CAST(ob);
849+
Py_VISIT(self->dict);
850+
return 0;
851+
}
852+
853+
static int
854+
ZstdDecompressor_clear(PyObject *ob)
855+
{
856+
ZstdDecompressor *self = ZstdDecompressor_CAST(ob);
857+
Py_CLEAR(self->dict);
858+
Py_CLEAR(self->unused_data);
859+
return 0;
860+
}
861+
840862
static PyType_Slot ZstdDecompressor_slots[] = {
841863
{Py_tp_new, _zstd_ZstdDecompressor_new},
842864
{Py_tp_dealloc, ZstdDecompressor_dealloc},
@@ -845,12 +867,14 @@ static PyType_Slot ZstdDecompressor_slots[] = {
845867
{Py_tp_members, ZstdDecompressor_members},
846868
{Py_tp_getset, ZstdDecompressor_getset},
847869
{Py_tp_doc, (char*)_zstd_ZstdDecompressor___init____doc__},
870+
{Py_tp_traverse, ZstdDecompressor_traverse},
871+
{Py_tp_clear, ZstdDecompressor_clear},
848872
{0}
849873
};
850874

851875
PyType_Spec ZstdDecompressor_type_spec = {
852876
.name = "_zstd.ZstdDecompressor",
853877
.basicsize = sizeof(ZstdDecompressor),
854-
.flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE,
878+
.flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC,
855879
.slots = ZstdDecompressor_slots,
856880
};

‎Modules/_zstd/zdict.c‎

Lines changed: 27 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -39,25 +39,25 @@ _zstd_ZstdDict_new(PyTypeObject *type, PyObject *Py_UNUSED(args), PyObject *Py_U
3939
return (PyObject*)self;
4040

4141
error:
42-
Py_XDECREF(self);
42+
PyObject_GC_UnTrack(self);
4343
return NULL;
4444
}
4545

4646
static void
4747
ZstdDict_dealloc(PyObject *ob)
4848
{
4949
ZstdDict *self = ZstdDict_CAST(ob);
50-
/* Free ZSTD_CDict instances */
51-
Py_XDECREF(self->c_dicts);
50+
51+
PyObject_GC_UnTrack(self);
5252

5353
/* Free ZSTD_DDict instance */
5454
ZSTD_freeDDict(self->d_dict);
5555

5656
/* Release dict_content after Free ZSTD_CDict/ZSTD_DDict instances */
57-
Py_XDECREF(self->dict_content);
57+
Py_CLEAR(self->dict_content);
5858

5959
PyTypeObject *tp = Py_TYPE(self);
60-
tp->tp_free((PyObject*)self);
60+
PyObject_GC_Del(ob);
6161
Py_DECREF(tp);
6262
}
6363

@@ -123,6 +123,8 @@ _zstd_ZstdDict___init___impl(ZstdDict *self, PyObject *dict_content,
123123
return -1;
124124
}
125125

126+
// Can only track self once self->dict_content is included
127+
PyObject_GC_Track(self);
126128
return 0;
127129
}
128130

@@ -238,6 +240,23 @@ ZstdDict_length(PyObject *ob)
238240
return Py_SIZE(self->dict_content);
239241
}
240242

243+
static int
244+
ZstdDict_traverse(PyObject *ob, visitproc visit, void *arg)
245+
{
246+
ZstdDict *self = ZstdDict_CAST(ob);
247+
Py_VISIT(self->c_dicts);
248+
Py_VISIT(self->dict_content);
249+
return 0;
250+
}
251+
252+
static int
253+
ZstdDict_clear(PyObject *ob)
254+
{
255+
ZstdDict *self = ZstdDict_CAST(ob);
256+
Py_CLEAR(self->dict_content);
257+
return 0;
258+
}
259+
241260
static PyType_Slot zstddict_slots[] = {
242261
{Py_tp_members, ZstdDict_members},
243262
{Py_tp_getset, ZstdDict_getset},
@@ -247,13 +266,14 @@ static PyType_Slot zstddict_slots[] = {
247266
{Py_tp_str, ZstdDict_str},
248267
{Py_tp_doc, (char*)_zstd_ZstdDict___init____doc__},
249268
{Py_sq_length, ZstdDict_length},
269+
{Py_tp_traverse, ZstdDict_traverse},
270+
{Py_tp_clear, ZstdDict_clear},
250271
{0}
251272
};
252273

253-
// TODO(emmatyping): ZstdDict (and others?) should become a GC object.
254274
PyType_Spec zstddict_type_spec = {
255275
.name = "_zstd.ZstdDict",
256276
.basicsize = sizeof(ZstdDict),
257-
.flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE,
277+
.flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC,
258278
.slots = zstddict_slots,
259279
};

0 commit comments

Comments
 (0)