@@ -178,8 +178,7 @@ _PyUnicodeWriter_InitWithBuffer(_PyUnicodeWriter *writer, PyObject *buffer)
178178 assert (PyUnstable_Object_IsUniquelyReferenced (buffer ));
179179
180180 memset (writer , 0 , sizeof (* writer ));
181- writer -> buffer = buffer ;
182- _PyUnicodeWriter_Update (writer );
181+ _PyUnicodeWriter_SetBuffer (writer , buffer );
183182 writer -> min_length = writer -> size ;
184183 assert (_PyUnicodeWriter_CanWrite (writer ));
185184}
@@ -204,16 +203,19 @@ _PyUnicodeWriter_PrepareInternal(_PyUnicodeWriter *writer,
204203
205204 maxchar = Py_MAX (maxchar , writer -> min_char );
206205
207- PyObject * newbuffer ;
206+ PyObject * new_buffer ;
208207 if (writer -> buffer == NULL ) {
209208 assert (!writer -> readonly );
209+
210210 // Do not overallocate at the first allocation, but use min_length
211- if (alloc < writer -> min_length )
211+ if (alloc < writer -> min_length ) {
212212 alloc = writer -> min_length ;
213+ }
213214
214- writer -> buffer = PyUnicode_New (alloc , maxchar );
215- if (writer -> buffer == NULL )
215+ new_buffer = PyUnicode_New (alloc , maxchar );
216+ if (new_buffer == NULL ) {
216217 return -1 ;
218+ }
217219 }
218220 else if (alloc > writer -> size ) {
219221 // Do not overallocate at the first allocation, but use min_length
@@ -223,38 +225,42 @@ _PyUnicodeWriter_PrepareInternal(_PyUnicodeWriter *writer,
223225 /* overallocate to limit the number of realloc() */
224226 alloc += alloc / OVERALLOCATE_FACTOR ;
225227 }
226- if (alloc < writer -> min_length )
228+ if (alloc < writer -> min_length ) {
227229 alloc = writer -> min_length ;
230+ }
228231
229232 if (maxchar > writer -> maxchar || writer -> readonly ) {
230233 /* resize + widen */
231234 maxchar = Py_MAX (maxchar , writer -> maxchar );
232- newbuffer = PyUnicode_New (alloc , maxchar );
233- if (newbuffer == NULL )
235+ new_buffer = PyUnicode_New (alloc , maxchar );
236+ if (new_buffer == NULL ) {
234237 return -1 ;
235- _PyUnicode_FastCopyCharacters (newbuffer , 0 ,
238+ }
239+ _PyUnicode_FastCopyCharacters (new_buffer , 0 ,
236240 writer -> buffer , 0 , writer -> pos );
237- writer -> readonly = 0 ;
238- Py_DECREF (writer -> buffer );
239- writer -> buffer = newbuffer ;
240241 }
241242 else {
242- newbuffer = _PyUnicode_ResizeCompact (writer -> buffer , alloc );
243- if (newbuffer == NULL )
243+ new_buffer = _PyUnicode_ResizeCompact (writer -> buffer , alloc );
244+ if (new_buffer == NULL ) {
244245 return -1 ;
245- writer -> buffer = newbuffer ;
246+ }
247+ // Do not DECREF the old buffer
248+ writer -> buffer = NULL ;
246249 }
247250 }
248- else if (maxchar > writer -> maxchar ) {
251+ else {
252+ assert (maxchar > writer -> maxchar );
249253 assert (!writer -> readonly );
250- newbuffer = PyUnicode_New (writer -> size , maxchar );
251- if (newbuffer == NULL )
254+
255+ new_buffer = PyUnicode_New (writer -> size , maxchar );
256+ if (new_buffer == NULL ) {
252257 return -1 ;
253- _PyUnicode_FastCopyCharacters (newbuffer , 0 ,
258+ }
259+ _PyUnicode_FastCopyCharacters (new_buffer , 0 ,
254260 writer -> buffer , 0 , writer -> pos );
255- Py_SETREF (writer -> buffer , newbuffer );
256261 }
257- _PyUnicodeWriter_Update (writer );
262+
263+ _PyUnicodeWriter_SetBuffer (writer , new_buffer );
258264 return 0 ;
259265
260266#undef OVERALLOCATE_FACTOR
@@ -316,11 +322,7 @@ _PyUnicodeWriter_WriteStr(_PyUnicodeWriter *writer, PyObject *str)
316322 if (maxchar > writer -> maxchar || len > writer -> size - writer -> pos ) {
317323 if (writer -> buffer == NULL && PyUnicode_CheckExact (str )) {
318324 assert (_PyUnicode_CheckConsistency (str , 1 ));
319- writer -> readonly = 1 ;
320- writer -> buffer = Py_NewRef (str );
321- _PyUnicodeWriter_Update (writer );
322- writer -> pos += len ;
323- // The next write will create a new buffer and copy the string
325+ _PyUnicodeWriter_SetReadOnly (writer , Py_NewRef (str ), len );
324326 return 0 ;
325327 }
326328 if (_PyUnicodeWriter_PrepareInternal (writer , len , maxchar ) == -1 )
@@ -457,10 +459,7 @@ _PyUnicodeWriter_WriteASCIIString(_PyUnicodeWriter *writer,
457459 if (str == NULL )
458460 return -1 ;
459461
460- writer -> readonly = 1 ;
461- writer -> buffer = str ;
462- _PyUnicodeWriter_Update (writer );
463- writer -> pos += len ;
462+ _PyUnicodeWriter_SetReadOnly (writer , str , len );
464463 return 0 ;
465464 }
466465
@@ -639,6 +638,26 @@ _PyUnicodeWriter_Finish(_PyUnicodeWriter *writer)
639638}
640639
641640
641+ PyObject *
642+ _PyUnicodeWriter_FinishWithSize (_PyUnicodeWriter * writer , Py_ssize_t size )
643+ {
644+ assert (0 <= size );
645+ if (writer -> buffer != NULL ) {
646+ assert (size <= writer -> pos );
647+ assert (size <= PyUnicode_GET_LENGTH (writer -> buffer ));
648+ if (size < writer -> pos ) {
649+ // Truncate the string: we may need to adjust the string kind
650+ writer -> recheck_maxchar = 1 ;
651+ }
652+ }
653+ else {
654+ assert (size == 0 );
655+ }
656+ writer -> pos = size ;
657+ return _PyUnicodeWriter_Finish (writer );
658+ }
659+
660+
642661PyObject *
643662PyUnicodeWriter_Finish (PyUnicodeWriter * writer )
644663{
0 commit comments