@@ -131,6 +131,127 @@ tuple_fromarray(PyObject* Py_UNUSED(module), PyObject *args)
131131}
132132
133133
134+ // --- PyTupleWriter type ---------------------------------------------------
135+
136+ typedef struct {
137+ PyObject_HEAD
138+ PyTupleWriter * writer ;
139+ } WriterObject ;
140+
141+
142+ static PyObject *
143+ writer_new (PyTypeObject * type , PyObject * args , PyObject * kwargs )
144+ {
145+ WriterObject * self = (WriterObject * )type -> tp_alloc (type , 0 );
146+ if (!self ) {
147+ return NULL ;
148+ }
149+ self -> writer = NULL ;
150+ return (PyObject * )self ;
151+ }
152+
153+
154+ static int
155+ writer_init (PyObject * self_raw , PyObject * args , PyObject * kwargs )
156+ {
157+ if (kwargs && PyDict_GET_SIZE (kwargs )) {
158+ PyErr_Format (PyExc_TypeError ,
159+ "PyTupleWriter() takes exactly no keyword arguments" );
160+ return -1 ;
161+ }
162+
163+ Py_ssize_t size ;
164+ if (!PyArg_ParseTuple (args , "n" , & size )) {
165+ return -1 ;
166+ }
167+
168+ WriterObject * self = (WriterObject * )self_raw ;
169+ if (self -> writer ) {
170+ PyTupleWriter_Discard (self -> writer );
171+ }
172+ self -> writer = PyTupleWriter_Create (size );
173+ if (self -> writer == NULL ) {
174+ return -1 ;
175+ }
176+ return 0 ;
177+ }
178+
179+
180+ static void
181+ writer_dealloc (PyObject * self_raw )
182+ {
183+ WriterObject * self = (WriterObject * )self_raw ;
184+ PyTypeObject * tp = Py_TYPE (self );
185+ if (self -> writer ) {
186+ PyTupleWriter_Discard (self -> writer );
187+ }
188+ tp -> tp_free (self );
189+ Py_DECREF (tp );
190+ }
191+
192+
193+ static inline int
194+ writer_check (WriterObject * self )
195+ {
196+ if (self -> writer == NULL ) {
197+ PyErr_SetString (PyExc_ValueError , "operation on finished writer" );
198+ return -1 ;
199+ }
200+ return 0 ;
201+ }
202+
203+
204+ static PyObject *
205+ writer_add (PyObject * self_raw , PyObject * item )
206+ {
207+ WriterObject * self = (WriterObject * )self_raw ;
208+ if (writer_check (self ) < 0 ) {
209+ return NULL ;
210+ }
211+
212+ if (PyTupleWriter_Add (self -> writer , item ) < 0 ) {
213+ return NULL ;
214+ }
215+ Py_RETURN_NONE ;
216+ }
217+
218+
219+ static PyObject *
220+ writer_finish (PyObject * self_raw , PyObject * Py_UNUSED (args ))
221+ {
222+ WriterObject * self = (WriterObject * )self_raw ;
223+ if (writer_check (self ) < 0 ) {
224+ return NULL ;
225+ }
226+
227+ PyObject * tuple = PyTupleWriter_Finish (self -> writer );
228+ self -> writer = NULL ;
229+ return tuple ;
230+ }
231+
232+
233+ static PyMethodDef writer_methods [] = {
234+ {"add" , _PyCFunction_CAST (writer_add ), METH_O },
235+ {"finish" , _PyCFunction_CAST (writer_finish ), METH_NOARGS },
236+ {NULL , NULL } /* sentinel */
237+ };
238+
239+ static PyType_Slot Writer_Type_slots [] = {
240+ {Py_tp_new , writer_new },
241+ {Py_tp_init , writer_init },
242+ {Py_tp_dealloc , writer_dealloc },
243+ {Py_tp_methods , writer_methods },
244+ {0 , 0 }, /* sentinel */
245+ };
246+
247+ static PyType_Spec Writer_spec = {
248+ .name = "_testcapi.PyTupleWriter" ,
249+ .basicsize = sizeof (WriterObject ),
250+ .flags = Py_TPFLAGS_DEFAULT ,
251+ .slots = Writer_Type_slots ,
252+ };
253+
254+
134255static PyMethodDef test_methods [] = {
135256 {"tuple_get_size" , tuple_get_size , METH_O },
136257 {"tuple_get_item" , tuple_get_item , METH_VARARGS },
@@ -148,5 +269,15 @@ _PyTestCapi_Init_Tuple(PyObject *m)
148269 return -1 ;
149270 }
150271
272+ PyTypeObject * writer_type = (PyTypeObject * )PyType_FromSpec (& Writer_spec );
273+ if (writer_type == NULL ) {
274+ return -1 ;
275+ }
276+ if (PyModule_AddType (m , writer_type ) < 0 ) {
277+ Py_DECREF (writer_type );
278+ return -1 ;
279+ }
280+ Py_DECREF (writer_type );
281+
151282 return 0 ;
152283}
0 commit comments